From 79524de91b9b8fc3bc73e8f40b4f1e0d3be4fb0b Mon Sep 17 00:00:00 2001 From: Ricardo Temperini <29879569+rtemperini@users.noreply.github.com> Date: Tue, 25 Aug 2026 18:42:54 +0200 Subject: [PATCH 01/13] feat(adk): promote allowlisted caller context onto every agent span Agent spans record what an agent did but not who asked for it, so traces cannot be filtered or grouped by the calling user, Slack thread, or ticket that triggered them. Read an operator-defined allowlist of context keys from W3C baggage and A2A message metadata and merge them into the request-scoped attribute bag. The existing span processor then stamps them on every span of the request, which is what trace-level filtering in Langfuse and comparable backends requires; attaching them to the root span alone leaves most views unfilterable. Baggage is the primary source because it already survives the controller, agent, sub-agent, and tool hops under the composite propagator, so a value set once at the edge needs no further plumbing. A2A message metadata covers callers that cannot set headers and, being per-message, takes precedence. Caller data is untrusted, so only allowlisted keys are read, the allowlist is capped, values are truncated and stripped of control characters, and every attribute is namespaced under kagent.context. so it cannot shadow a semantic convention attribute such as service.name. The allowlist is empty by default, which disables promotion entirely. Signed-off-by: Ricardo Temperini <29879569+rtemperini@users.noreply.github.com> Co-authored-by: Cursor --- go/adk/pkg/a2a/executor.go | 4 + go/adk/pkg/telemetry/context_attributes.go | 140 ++++++++++ .../pkg/telemetry/context_attributes_test.go | 250 ++++++++++++++++++ 3 files changed, 394 insertions(+) create mode 100644 go/adk/pkg/telemetry/context_attributes.go create mode 100644 go/adk/pkg/telemetry/context_attributes_test.go diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index 57d3120079..2201806485 100644 --- a/go/adk/pkg/a2a/executor.go +++ b/go/adk/pkg/a2a/executor.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "iter" + "maps" "strings" a2atype "github.com/a2aproject/a2a-go/v2/a2a" @@ -129,6 +130,9 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.ExecutorCon if e.appName != "" { spanAttributes["kagent.app_name"] = e.appName } + // Allowlisted caller context joins the request-scoped bag rather than a + // single span, so tool, sub-agent, and model spans all carry it. + maps.Copy(spanAttributes, telemetry.CallerContextAttributes(ctx, reqCtx.Message.Metadata)) ctx = telemetry.SetKAgentSpanAttributes(ctx, spanAttributes) ctx, invocationSpan := telemetry.StartInvocationSpan(ctx) defer invocationSpan.End() diff --git a/go/adk/pkg/telemetry/context_attributes.go b/go/adk/pkg/telemetry/context_attributes.go new file mode 100644 index 0000000000..f0de636101 --- /dev/null +++ b/go/adk/pkg/telemetry/context_attributes.go @@ -0,0 +1,140 @@ +package telemetry + +import ( + "context" + "os" + "strconv" + "strings" + "unicode" + + "go.opentelemetry.io/otel/baggage" +) + +const ( + // traceContextKeysEnvVar holds a comma-separated allowlist of caller-supplied + // context keys to promote onto agent spans. Unset or empty (the default) + // disables promotion entirely. + traceContextKeysEnvVar = "KAGENT_TRACE_CONTEXT_KEYS" + + // contextAttributePrefix namespaces every promoted value. Because the prefix + // is applied unconditionally, caller-supplied data cannot shadow a semantic + // convention attribute such as service.name. + contextAttributePrefix = "kagent.context." + + maxContextKeys = 32 + maxContextKeyLength = 64 + maxContextValueLength = 256 +) + +// CallerContextAttributes returns the caller-supplied context values that an +// operator allowlisted through KAGENT_TRACE_CONTEXT_KEYS. Merge the result into +// the request-scoped attribute bag (see SetKAgentSpanAttributes) so every span +// of the request carries them: trace-level filtering in backends such as +// Langfuse matches on each span, not only on the root. +// +// Values are read from W3C baggage first and then from the A2A message +// metadata, which is the more specific source for a single message and +// therefore wins. Both are untrusted input, so keys must appear in the +// allowlist, values are stripped of control characters and truncated, and every +// attribute is namespaced under contextAttributePrefix. +// +// Returns nil when the allowlist is empty, which is the default. +func CallerContextAttributes(ctx context.Context, metadata map[string]any) map[string]string { + keys := allowedContextKeys() + if len(keys) == 0 { + return nil + } + + bag := baggage.FromContext(ctx) + attrs := make(map[string]string, len(keys)) + for _, key := range keys { + value := sanitizeContextValue(bag.Member(key).Value()) + if scalar, ok := scalarString(metadata[key]); ok { + value = sanitizeContextValue(scalar) + } + if value == "" { + continue + } + attrs[contextAttributePrefix+key] = value + } + if len(attrs) == 0 { + return nil + } + return attrs +} + +// allowedContextKeys parses the KAGENT_TRACE_CONTEXT_KEYS allowlist. Keys that +// are empty, over-long, or contain whitespace or control characters are +// dropped, and the list is capped at maxContextKeys so a misconfigured +// allowlist cannot inflate span cardinality without bound. +func allowedContextKeys() []string { + raw := strings.TrimSpace(os.Getenv(traceContextKeysEnvVar)) + if raw == "" { + return nil + } + + keys := make([]string, 0, maxContextKeys) + seen := make(map[string]struct{}, maxContextKeys) + for key := range strings.SplitSeq(raw, ",") { + key = strings.TrimSpace(key) + if key == "" || len(key) > maxContextKeyLength || !isAttributeKey(key) { + continue + } + if _, duplicate := seen[key]; duplicate { + continue + } + seen[key] = struct{}{} + keys = append(keys, key) + if len(keys) == maxContextKeys { + break + } + } + return keys +} + +// isAttributeKey reports whether key is safe to use as a span attribute name. +func isAttributeKey(key string) bool { + return strings.IndexFunc(key, func(r rune) bool { + return unicode.IsControl(r) || unicode.IsSpace(r) + }) < 0 +} + +// scalarString renders a JSON scalar from A2A message metadata as a string. +// Objects and arrays are skipped: they are unbounded in size and carry no +// useful meaning as a span attribute value. +func scalarString(value any) (string, bool) { + switch v := value.(type) { + case string: + return v, true + case bool: + return strconv.FormatBool(v), true + case float64: + return strconv.FormatFloat(v, 'g', -1, 64), true + case int: + return strconv.Itoa(v), true + case int64: + return strconv.FormatInt(v, 10), true + default: + return "", false + } +} + +// sanitizeContextValue makes an untrusted value safe to attach to a span: +// control characters are dropped so a value cannot forge structure in a +// downstream trace or log renderer, and the result is truncated to bound the +// size of exported spans. +func sanitizeContextValue(value string) string { + cleaned := strings.Map(func(r rune) rune { + if unicode.IsControl(r) { + return -1 + } + return r + }, value) + cleaned = strings.TrimSpace(cleaned) + // Truncate by rune, not byte, so the limit means the same thing here as it + // does in the Python runtime. + if runes := []rune(cleaned); len(runes) > maxContextValueLength { + cleaned = string(runes[:maxContextValueLength]) + } + return cleaned +} diff --git a/go/adk/pkg/telemetry/context_attributes_test.go b/go/adk/pkg/telemetry/context_attributes_test.go new file mode 100644 index 0000000000..5a4eb06c6c --- /dev/null +++ b/go/adk/pkg/telemetry/context_attributes_test.go @@ -0,0 +1,250 @@ +package telemetry + +import ( + "context" + "strconv" + "strings" + "testing" + + "go.opentelemetry.io/otel/baggage" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" +) + +func baggageContext(t *testing.T, members map[string]string) context.Context { + t.Helper() + + built := make([]baggage.Member, 0, len(members)) + for key, value := range members { + member, err := baggage.NewMember(key, value) + if err != nil { + t.Fatalf("baggage.NewMember(%q, %q): %v", key, value, err) + } + built = append(built, member) + } + bag, err := baggage.New(built...) + if err != nil { + t.Fatalf("baggage.New: %v", err) + } + return baggage.ContextWithBaggage(context.Background(), bag) +} + +func TestCallerContextAttributes(t *testing.T) { + tests := []struct { + name string + allowlist string + baggageVals map[string]string + metadata map[string]any + want map[string]string + }{ + { + name: "disabled by default", + allowlist: "", + baggageVals: map[string]string{"user.email": "ada@example.com"}, + metadata: map[string]any{"thread_id": "T123"}, + want: nil, + }, + { + name: "promotes allowlisted baggage", + allowlist: "user.email,user.name", + baggageVals: map[string]string{"user.email": "ada@example.com", "user.name": "Ada"}, + want: map[string]string{ + "kagent.context.user.email": "ada@example.com", + "kagent.context.user.name": "Ada", + }, + }, + { + name: "promotes allowlisted message metadata", + allowlist: "thread_id,channel", + metadata: map[string]any{"thread_id": "1717171.4242", "channel": "C0AB1"}, + want: map[string]string{ + "kagent.context.thread_id": "1717171.4242", + "kagent.context.channel": "C0AB1", + }, + }, + { + name: "message metadata overrides baggage", + allowlist: "user.email", + baggageVals: map[string]string{"user.email": "from-baggage@example.com"}, + metadata: map[string]any{"user.email": "from-metadata@example.com"}, + want: map[string]string{"kagent.context.user.email": "from-metadata@example.com"}, + }, + { + name: "ignores keys outside the allowlist", + allowlist: "thread_id", + baggageVals: map[string]string{"secret.token": "s3cret"}, + metadata: map[string]any{"thread_id": "T1", "customer.pan": "4111111111111111"}, + want: map[string]string{"kagent.context.thread_id": "T1"}, + }, + { + name: "renders scalar metadata types", + allowlist: "count,ratio,enabled", + metadata: map[string]any{ + "count": int64(7), + "ratio": float64(2.5), + "enabled": true, + }, + want: map[string]string{ + "kagent.context.count": "7", + "kagent.context.ratio": "2.5", + "kagent.context.enabled": "true", + }, + }, + { + name: "skips non-scalar and empty metadata values", + allowlist: "nested,list,blank", + metadata: map[string]any{ + "nested": map[string]any{"a": "b"}, + "list": []string{"a"}, + "blank": "", + }, + want: nil, + }, + { + name: "strips control characters", + allowlist: "note", + metadata: map[string]any{"note": "line\nbreak\tand\x00nul"}, + want: map[string]string{"kagent.context.note": "linebreakandnul"}, + }, + { + name: "ignores allowlist entries that are not valid attribute keys", + allowlist: "good, bad key ,\tanother\tbad", + metadata: map[string]any{"good": "yes", "bad key": "no"}, + want: map[string]string{"kagent.context.good": "yes"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Setenv(traceContextKeysEnvVar, tt.allowlist) + + got := CallerContextAttributes(baggageContext(t, tt.baggageVals), tt.metadata) + + if len(got) != len(tt.want) { + t.Fatalf("got %v, want %v", got, tt.want) + } + for key, want := range tt.want { + if got[key] != want { + t.Errorf("%s = %q, want %q", key, got[key], want) + } + } + }) + } +} + +func TestCallerContextAttributes_TruncatesLongValues(t *testing.T) { + t.Setenv(traceContextKeysEnvVar, "note") + + got := CallerContextAttributes(context.Background(), map[string]any{ + "note": strings.Repeat("a", maxContextValueLength*2), + }) + + if len(got["kagent.context.note"]) != maxContextValueLength { + t.Errorf("value length = %d, want %d", len(got["kagent.context.note"]), maxContextValueLength) + } +} + +func TestAllowedContextKeys_CapsListLength(t *testing.T) { + keys := make([]string, 0, maxContextKeys*2) + for i := range maxContextKeys * 2 { + keys = append(keys, "key"+strconv.Itoa(i)) + } + t.Setenv(traceContextKeysEnvVar, strings.Join(keys, ",")) + + if got := len(allowedContextKeys()); got != maxContextKeys { + t.Errorf("allowlist length = %d, want %d", got, maxContextKeys) + } +} + +func TestAllowedContextKeys_DropsOverLongAndDuplicateKeys(t *testing.T) { + t.Setenv(traceContextKeysEnvVar, "a,a,"+strings.Repeat("b", maxContextKeyLength+1)+",c") + + got := allowedContextKeys() + + want := []string{"a", "c"} + if len(got) != len(want) { + t.Fatalf("got %v, want %v", got, want) + } + for i, key := range want { + if got[i] != key { + t.Errorf("key %d = %q, want %q", i, got[i], key) + } + } +} + +// Langfuse and comparable backends filter on attributes present on each span, +// so the promoted values must reach descendants, not just the root span. +func TestCallerContextAttributes_ReachEverySpan(t *testing.T) { + t.Setenv(traceContextKeysEnvVar, "user.email") + + exporter := tracetest.NewInMemoryExporter() + tp := sdktrace.NewTracerProvider( + sdktrace.WithSyncer(exporter), + sdktrace.WithSpanProcessor(kagentAttributesSpanProcessor{}), + ) + t.Cleanup(func() { + _ = tp.Shutdown(context.Background()) + }) + + ctx := baggageContext(t, map[string]string{"user.email": "ada@example.com"}) + ctx = SetKAgentSpanAttributes(ctx, CallerContextAttributes(ctx, nil)) + + tracer := tp.Tracer("test") + ctx, root := tracer.Start(ctx, "root") + ctx, tool := tracer.Start(ctx, "execute_tool") + _, model := tracer.Start(ctx, "generate_content") + model.End() + tool.End() + root.End() + + spans := exporter.GetSpans() + if len(spans) != 3 { + t.Fatalf("expected 3 spans, got %d", len(spans)) + } + for _, name := range []string{"root", "execute_tool", "generate_content"} { + attrs := spanAttributesByName(t, spans, name) + if got := attrs["kagent.context.user.email"].AsString(); got != "ada@example.com" { + t.Errorf("span %q: kagent.context.user.email = %q, want %q", name, got, "ada@example.com") + } + } +} + +func TestCallerContextAttributes_DisabledLeavesSpansUnchanged(t *testing.T) { + t.Setenv(traceContextKeysEnvVar, "") + + exporter := tracetest.NewInMemoryExporter() + tp := sdktrace.NewTracerProvider( + sdktrace.WithSyncer(exporter), + sdktrace.WithSpanProcessor(kagentAttributesSpanProcessor{}), + ) + t.Cleanup(func() { + _ = tp.Shutdown(context.Background()) + }) + + ctx := baggageContext(t, map[string]string{"user.email": "ada@example.com"}) + ctx = SetKAgentSpanAttributes(ctx, CallerContextAttributes(ctx, map[string]any{"thread_id": "T1"})) + + _, span := tp.Tracer("test").Start(ctx, "root") + span.End() + + for _, attr := range exporter.GetSpans()[0].Attributes { + if strings.HasPrefix(string(attr.Key), contextAttributePrefix) { + t.Errorf("unexpected promoted attribute %q", attr.Key) + } + } +} + +// The kagent.context. prefix is what makes caller-supplied data unable to +// shadow a semantic convention attribute, even if an operator allowlists one. +func TestCallerContextAttributes_CannotShadowSemanticConventions(t *testing.T) { + t.Setenv(traceContextKeysEnvVar, "service.name") + + got := CallerContextAttributes(context.Background(), map[string]any{"service.name": "impostor"}) + + if _, shadowed := got["service.name"]; shadowed { + t.Error("service.name must not be settable by a caller") + } + if got["kagent.context.service.name"] != "impostor" { + t.Errorf("got %v, want the value namespaced under %q", got, contextAttributePrefix) + } +} From ba8d206224b2a53bf586f23c2f08fd41fac1fb35 Mon Sep 17 00:00:00 2001 From: Ricardo Temperini <29879569+rtemperini@users.noreply.github.com> Date: Tue, 25 Aug 2026 18:43:02 +0200 Subject: [PATCH 02/13] feat(python): promote allowlisted caller context onto every agent span Bring the Python runtimes to parity with the Go ADK. The Go runtime already reads A2A message metadata into span attributes (#1734, #1737); the Python runtimes ignored inbound metadata entirely, so agents on the Python runtime had no way to get caller identity onto their traces. Mirror the Go implementation exactly: the same KAGENT_TRACE_CONTEXT_KEYS allowlist, the same baggage-then-metadata precedence, the same limits, the same control character stripping, and the same kagent.context. namespace, so the two runtimes cannot drift. Values are merged into the request-scoped attribute bag that KagentAttributesSpanProcessor stamps onto every span. The ADK, LangGraph, and CrewAI executors all promote context through the shared helper. Reading a Message's protobuf Struct metadata moves into kagent.core.a2a.read_message_metadata rather than being repeated per package. Also assert that the OTel SDK's default propagator carries baggage: the Python runtime relies on that default rather than configuring a propagator, so an SDK change that dropped it would silently break the baggage path. Signed-off-by: Ricardo Temperini <29879569+rtemperini@users.noreply.github.com> Co-authored-by: Cursor --- .../src/kagent/adk/_agent_executor.py | 6 + .../src/kagent/core/a2a/__init__.py | 2 + .../src/kagent/core/a2a/_consts.py | 19 ++ .../src/kagent/core/tracing/__init__.py | 3 +- .../core/tracing/_context_attributes.py | 131 +++++++++++++ .../tests/test_caller_context_attributes.py | 175 ++++++++++++++++++ .../tests/test_read_metadata_value.py | 27 ++- .../tests/test_tracing_configure.py | 14 ++ .../src/kagent/crewai/_executor.py | 7 +- .../src/kagent/langgraph/_executor.py | 6 + 10 files changed, 387 insertions(+), 3 deletions(-) create mode 100644 python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py create mode 100644 python/packages/kagent-core/tests/test_caller_context_attributes.py diff --git a/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py b/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py index 18107e1ca5..fb8cf98c5f 100644 --- a/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py +++ b/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py @@ -33,7 +33,9 @@ get_kagent_metadata_key, hitl_activated, now_timestamp, + read_message_metadata, ) +from kagent.core.tracing import caller_context_attributes from kagent.core.tracing._span_processor import clear_kagent_span_attributes, set_kagent_span_attributes from pydantic import BaseModel @@ -131,6 +133,10 @@ async def execute(self, context: RequestContext, event_queue: EventQueue) -> Non "gen_ai.task.id": context.task_id, "gen_ai.conversation.id": run_request.session_id, } + # Allowlisted caller context joins the request-scoped bag rather + # than a single span, so tool, sub-agent, and model spans all + # carry it. + span_attributes.update(caller_context_attributes(read_message_metadata(context.message))) context_token = set_kagent_span_attributes( {key: value for key, value in span_attributes.items() if value is not None} ) diff --git a/python/packages/kagent-core/src/kagent/core/a2a/__init__.py b/python/packages/kagent-core/src/kagent/core/a2a/__init__.py index 73cc2a6504..4fafa54095 100644 --- a/python/packages/kagent-core/src/kagent/core/a2a/__init__.py +++ b/python/packages/kagent-core/src/kagent/core/a2a/__init__.py @@ -8,6 +8,7 @@ A2A_DATA_PART_METADATA_TYPE_KEY, ADK_METADATA_KEY_PREFIX, get_kagent_metadata_key, + read_message_metadata, read_metadata_value, ) from ._context import get_request_user_id, set_request_user_id @@ -49,6 +50,7 @@ "KAgentGrpcServerCallContextBuilder", "now_timestamp", "get_kagent_metadata_key", + "read_message_metadata", "read_metadata_value", "ADK_METADATA_KEY_PREFIX", "A2A_DATA_PART_METADATA_TYPE_KEY", diff --git a/python/packages/kagent-core/src/kagent/core/a2a/_consts.py b/python/packages/kagent-core/src/kagent/core/a2a/_consts.py index 74608dbd82..568d71d192 100644 --- a/python/packages/kagent-core/src/kagent/core/a2a/_consts.py +++ b/python/packages/kagent-core/src/kagent/core/a2a/_consts.py @@ -1,3 +1,8 @@ +from typing import Any, Optional + +from a2a.types import Message +from google.protobuf.json_format import MessageToDict + # A2A DataPart metadata constants. # These values MUST match the upstream google-adk definitions in # google.adk.a2a.converters.part_converter. A sync-check test in @@ -30,6 +35,20 @@ def get_kagent_metadata_key(key: str) -> str: return f"{KAGENT_METADATA_KEY_PREFIX}{key}" +def read_message_metadata(message: Optional[Message]) -> dict[str, Any]: + """Return a Message's protobuf ``Struct`` metadata as a plain dict. + + Args: + message: The A2A message to read (may be ``None``). + + Returns: + The decoded metadata, or an empty dict when the message carries none. + """ + if message is None or not message.HasField("metadata"): + return {} + return MessageToDict(message.metadata) + + def read_metadata_value(metadata: dict | None, key: str, default=None): """Read a metadata value, checking ``adk_`` first then ``kagent_``. diff --git a/python/packages/kagent-core/src/kagent/core/tracing/__init__.py b/python/packages/kagent-core/src/kagent/core/tracing/__init__.py index 826775371d..87307f1e71 100644 --- a/python/packages/kagent-core/src/kagent/core/tracing/__init__.py +++ b/python/packages/kagent-core/src/kagent/core/tracing/__init__.py @@ -1,3 +1,4 @@ +from ._context_attributes import caller_context_attributes from ._utils import configure, force_flush -__all__ = ["configure", "force_flush"] +__all__ = ["caller_context_attributes", "configure", "force_flush"] diff --git a/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py b/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py new file mode 100644 index 0000000000..2956217367 --- /dev/null +++ b/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py @@ -0,0 +1,131 @@ +"""Promote allowlisted caller context onto every span of an agent request.""" + +import os +from typing import Any, Optional + +from opentelemetry import baggage +from opentelemetry import context as otel_context + +# Comma-separated allowlist of caller-supplied context keys to promote onto +# agent spans. Unset or empty (the default) disables promotion entirely. +TRACE_CONTEXT_KEYS_ENV_VAR = "KAGENT_TRACE_CONTEXT_KEYS" + +# Namespaces every promoted value. Because the prefix is applied +# unconditionally, caller-supplied data cannot shadow a semantic convention +# attribute such as ``service.name``. +CONTEXT_ATTRIBUTE_PREFIX = "kagent.context." + +MAX_CONTEXT_KEYS = 32 +MAX_CONTEXT_KEY_LENGTH = 64 +MAX_CONTEXT_VALUE_LENGTH = 256 + + +def caller_context_attributes( + metadata: Optional[dict[str, Any]] = None, + context: Optional[otel_context.Context] = None, +) -> dict[str, str]: + """Return the caller context values an operator allowlisted for tracing. + + Merge the result into the request-scoped attribute bag (see + ``set_kagent_span_attributes``) so every span of the request carries the + values: trace-level filtering in backends such as Langfuse matches on each + span, not only on the root. + + Values are read from W3C baggage first and then from the A2A message + metadata, which is the more specific source for a single message and + therefore wins. Both are untrusted input, so keys must appear in the + allowlist, values are stripped of control characters and truncated, and + every attribute is namespaced under ``CONTEXT_ATTRIBUTE_PREFIX``. + + Args: + metadata: A2A message metadata as a plain dict (may be ``None``). + context: OTel context to read baggage from. Defaults to the current one. + + Returns: + Prefixed attribute name to sanitised value. Empty when the allowlist is + empty, which is the default. + """ + keys = _allowed_context_keys() + if not keys: + return {} + + bag = baggage.get_all(context) + attributes: dict[str, str] = {} + for key in keys: + value = _sanitize_context_value(bag.get(key)) + if metadata is not None: + scalar = _scalar_string(metadata.get(key)) + if scalar is not None: + value = _sanitize_context_value(scalar) + if not value: + continue + attributes[CONTEXT_ATTRIBUTE_PREFIX + key] = value + return attributes + + +def _allowed_context_keys() -> list[str]: + """Parse the ``KAGENT_TRACE_CONTEXT_KEYS`` allowlist. + + Keys that are empty, over-long, or contain whitespace or control characters + are dropped, and the list is capped at ``MAX_CONTEXT_KEYS`` so a + misconfigured allowlist cannot inflate span cardinality without bound. + """ + raw = os.getenv(TRACE_CONTEXT_KEYS_ENV_VAR, "").strip() + if not raw: + return [] + + keys: list[str] = [] + for candidate in raw.split(","): + key = candidate.strip() + if not key or len(key) > MAX_CONTEXT_KEY_LENGTH or not _is_attribute_key(key): + continue + if key in keys: + continue + keys.append(key) + if len(keys) == MAX_CONTEXT_KEYS: + break + return keys + + +def _is_attribute_key(key: str) -> bool: + """Report whether *key* is safe to use as a span attribute name.""" + return not any(_is_control(char) or char.isspace() for char in key) + + +def _is_control(char: str) -> bool: + """Match Go's ``unicode.IsControl`` so both runtimes sanitise identically.""" + code_point = ord(char) + return code_point < 0x20 or 0x7F <= code_point <= 0x9F + + +def _scalar_string(value: Any) -> Optional[str]: + """Render a JSON scalar from A2A message metadata as a string. + + Objects and arrays are skipped: they are unbounded in size and carry no + useful meaning as a span attribute value. ``None`` means "no scalar here", + which leaves any baggage value for the same key in place. + """ + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, int): + return str(value) + if isinstance(value, float): + # protobuf Struct has a single numeric type, so a JSON integer arrives + # as a float. Render it without the trailing ".0" to match the Go ADK. + return str(int(value)) if value.is_integer() else repr(value) + if isinstance(value, str): + return value + return None + + +def _sanitize_context_value(value: Any) -> str: + """Make an untrusted value safe to attach to a span. + + Control characters are dropped so a value cannot forge structure in a + downstream trace or log renderer, and the result is truncated to bound the + size of exported spans. + """ + if not isinstance(value, str): + return "" + cleaned = "".join(char for char in value if not _is_control(char)).strip() + return cleaned[:MAX_CONTEXT_VALUE_LENGTH] diff --git a/python/packages/kagent-core/tests/test_caller_context_attributes.py b/python/packages/kagent-core/tests/test_caller_context_attributes.py new file mode 100644 index 0000000000..c82b9311e9 --- /dev/null +++ b/python/packages/kagent-core/tests/test_caller_context_attributes.py @@ -0,0 +1,175 @@ +import pytest +from opentelemetry import baggage +from opentelemetry import context as otel_context +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from kagent.core.tracing import caller_context_attributes +from kagent.core.tracing._context_attributes import ( + CONTEXT_ATTRIBUTE_PREFIX, + MAX_CONTEXT_KEY_LENGTH, + MAX_CONTEXT_KEYS, + MAX_CONTEXT_VALUE_LENGTH, + TRACE_CONTEXT_KEYS_ENV_VAR, + _allowed_context_keys, +) +from kagent.core.tracing._span_processor import ( + KagentAttributesSpanProcessor, + clear_kagent_span_attributes, + set_kagent_span_attributes, +) + + +def baggage_context(members: dict[str, str]) -> otel_context.Context: + context = otel_context.Context() + for key, value in members.items(): + context = baggage.set_baggage(key, value, context) + return context + + +class TestCallerContextAttributes: + """Tests for allowlist-driven promotion of caller context onto spans.""" + + def test_disabled_by_default(self, monkeypatch): + monkeypatch.delenv(TRACE_CONTEXT_KEYS_ENV_VAR, raising=False) + assert ( + caller_context_attributes( + {"thread_id": "T1"}, + baggage_context({"user.email": "ada@example.com"}), + ) + == {} + ) + + def test_promotes_allowlisted_baggage(self, monkeypatch): + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "user.email,user.name") + assert caller_context_attributes( + None, baggage_context({"user.email": "ada@example.com", "user.name": "Ada"}) + ) == { + "kagent.context.user.email": "ada@example.com", + "kagent.context.user.name": "Ada", + } + + def test_promotes_allowlisted_message_metadata(self, monkeypatch): + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "thread_id,channel") + assert caller_context_attributes({"thread_id": "1717171.4242", "channel": "C0AB1"}) == { + "kagent.context.thread_id": "1717171.4242", + "kagent.context.channel": "C0AB1", + } + + def test_message_metadata_overrides_baggage(self, monkeypatch): + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "user.email") + assert caller_context_attributes( + {"user.email": "from-metadata@example.com"}, + baggage_context({"user.email": "from-baggage@example.com"}), + ) == {"kagent.context.user.email": "from-metadata@example.com"} + + def test_ignores_keys_outside_the_allowlist(self, monkeypatch): + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "thread_id") + assert caller_context_attributes( + {"thread_id": "T1", "customer.pan": "4111111111111111"}, + baggage_context({"secret.token": "s3cret"}), + ) == {"kagent.context.thread_id": "T1"} + + @pytest.mark.parametrize( + ("value", "expected"), + [ + (True, "true"), + (False, "false"), + (7, "7"), + # protobuf Struct has one numeric type, so JSON integers arrive as floats. + (3.0, "3"), + (2.5, "2.5"), + ], + ) + def test_renders_scalar_metadata_types(self, monkeypatch, value, expected): + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "value") + assert caller_context_attributes({"value": value}) == {"kagent.context.value": expected} + + @pytest.mark.parametrize("value", [{"a": "b"}, ["a"], "", None]) + def test_skips_non_scalar_and_empty_metadata_values(self, monkeypatch, value): + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "value") + assert caller_context_attributes({"value": value}) == {} + + def test_strips_control_characters(self, monkeypatch): + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "note") + assert caller_context_attributes({"note": "line\nbreak\tand\x00nul"}) == { + "kagent.context.note": "linebreakandnul" + } + + def test_truncates_long_values(self, monkeypatch): + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "note") + promoted = caller_context_attributes({"note": "a" * (MAX_CONTEXT_VALUE_LENGTH * 2)}) + assert len(promoted["kagent.context.note"]) == MAX_CONTEXT_VALUE_LENGTH + + def test_cannot_shadow_semantic_conventions(self, monkeypatch): + """The prefix is what stops caller data replacing service.name and friends.""" + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "service.name") + promoted = caller_context_attributes({"service.name": "impostor"}) + assert "service.name" not in promoted + assert promoted == {f"{CONTEXT_ATTRIBUTE_PREFIX}service.name": "impostor"} + + +class TestAllowedContextKeys: + """Tests for allowlist parsing and its bounds.""" + + def test_caps_list_length(self, monkeypatch): + keys = ",".join(f"key{index}" for index in range(MAX_CONTEXT_KEYS * 2)) + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, keys) + assert len(_allowed_context_keys()) == MAX_CONTEXT_KEYS + + def test_drops_over_long_and_duplicate_keys(self, monkeypatch): + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, f"a,a,{'b' * (MAX_CONTEXT_KEY_LENGTH + 1)},c") + assert _allowed_context_keys() == ["a", "c"] + + def test_drops_keys_that_are_not_valid_attribute_names(self, monkeypatch): + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "good, bad key ,\tanother\tbad") + assert _allowed_context_keys() == ["good"] + + def test_empty_allowlist_disables_promotion(self, monkeypatch): + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, " , ,") + assert _allowed_context_keys() == [] + + +class TestPromotedAttributesReachEverySpan: + """Langfuse and comparable backends filter on attributes present on each + span, so promoted values must reach descendants, not just the root span.""" + + @staticmethod + def record_spans(span_attributes: dict) -> dict[str, dict]: + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + provider.add_span_processor(KagentAttributesSpanProcessor()) + tracer = provider.get_tracer("test") + + token = set_kagent_span_attributes(span_attributes) + try: + with tracer.start_as_current_span("root"): + with tracer.start_as_current_span("execute_tool"): + with tracer.start_as_current_span("generate_content"): + pass + finally: + clear_kagent_span_attributes(token) + provider.shutdown() + + return {span.name: dict(span.attributes or {}) for span in exporter.get_finished_spans()} + + def test_flag_on_stamps_every_span(self, monkeypatch): + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "user.email") + promoted = caller_context_attributes(None, baggage_context({"user.email": "ada@example.com"})) + + spans = self.record_spans(promoted) + + assert set(spans) == {"root", "execute_tool", "generate_content"} + for attributes in spans.values(): + assert attributes["kagent.context.user.email"] == "ada@example.com" + + def test_flag_off_leaves_spans_unchanged(self, monkeypatch): + monkeypatch.delenv(TRACE_CONTEXT_KEYS_ENV_VAR, raising=False) + promoted = caller_context_attributes({"thread_id": "T1"}, baggage_context({"user.email": "ada@example.com"})) + + spans = self.record_spans(promoted) + + for attributes in spans.values(): + assert not [key for key in attributes if key.startswith(CONTEXT_ATTRIBUTE_PREFIX)] diff --git a/python/packages/kagent-core/tests/test_read_metadata_value.py b/python/packages/kagent-core/tests/test_read_metadata_value.py index 4abd6cbfa6..35112a4a0b 100644 --- a/python/packages/kagent-core/tests/test_read_metadata_value.py +++ b/python/packages/kagent-core/tests/test_read_metadata_value.py @@ -1,6 +1,31 @@ import pytest +from a2a.types import Message, Role -from kagent.core.a2a import read_metadata_value +from kagent.core.a2a import read_message_metadata, read_metadata_value + + +class TestReadMessageMetadata: + """Tests for decoding a Message's protobuf Struct metadata.""" + + def test_returns_empty_dict_for_none_message(self): + assert read_message_metadata(None) == {} + + def test_returns_empty_dict_when_metadata_unset(self): + assert read_message_metadata(Message(role=Role.ROLE_USER, message_id="m")) == {} + + def test_decodes_scalar_and_nested_values(self): + message = Message( + role=Role.ROLE_USER, + message_id="m", + metadata={"thread_id": "T1", "attempt": 3, "flags": {"dry_run": True}}, + ) + + assert read_message_metadata(message) == { + "thread_id": "T1", + # protobuf Struct stores every number as a double. + "attempt": 3.0, + "flags": {"dry_run": True}, + } class TestReadMetadataValue: diff --git a/python/packages/kagent-core/tests/test_tracing_configure.py b/python/packages/kagent-core/tests/test_tracing_configure.py index 3ad0b3e3f4..da9f0ac810 100644 --- a/python/packages/kagent-core/tests/test_tracing_configure.py +++ b/python/packages/kagent-core/tests/test_tracing_configure.py @@ -2,6 +2,7 @@ from types import SimpleNamespace import pytest +from opentelemetry.baggage import get_baggage from opentelemetry.propagate import get_global_textmap from opentelemetry.trace import get_current_span @@ -199,6 +200,19 @@ def test_otel_sdk_default_propagator_includes_w3c_tracecontext(): assert get_current_span(ctx).get_span_context().trace_id == trace_id +def test_otel_sdk_default_propagator_includes_baggage(): + """The OTEL SDK must propagate W3C Baggage by default. + + Baggage is how caller identity and context reach an agent and its + sub-agents (see caller_context_attributes). If an OTEL SDK upgrade drops + baggage from the default propagator, this test will fail and explicit + configuration will be needed. + """ + ctx = get_global_textmap().extract({"baggage": "user.email=ada%40example.com"}) + + assert get_baggage("user.email", ctx) == "ada@example.com" + + @pytest.mark.parametrize( ("signal", "env", "expected"), [ diff --git a/python/packages/kagent-crewai/src/kagent/crewai/_executor.py b/python/packages/kagent-crewai/src/kagent/crewai/_executor.py index 0010247726..dc5a80b103 100644 --- a/python/packages/kagent-crewai/src/kagent/crewai/_executor.py +++ b/python/packages/kagent-crewai/src/kagent/crewai/_executor.py @@ -22,7 +22,8 @@ TaskStatusUpdateEvent, ) from google.protobuf.json_format import MessageToDict -from kagent.core.a2a import get_kagent_metadata_key, now_timestamp +from kagent.core.a2a import get_kagent_metadata_key, now_timestamp, read_message_metadata +from kagent.core.tracing import caller_context_attributes from kagent.core.tracing._span_processor import ( clear_kagent_span_attributes, set_kagent_span_attributes, @@ -193,4 +194,8 @@ def _convert_a2a_request_to_span_attributes( if request.task_id: span_attributes["gen_ai.task.id"] = request.task_id + # Allowlisted caller context joins the request-scoped bag rather than a + # single span, so tool, sub-agent, and model spans all carry it. + span_attributes.update(caller_context_attributes(read_message_metadata(request.message))) + return span_attributes diff --git a/python/packages/kagent-langgraph/src/kagent/langgraph/_executor.py b/python/packages/kagent-langgraph/src/kagent/langgraph/_executor.py index f3ff8a6f8a..46c76ddc3a 100644 --- a/python/packages/kagent-langgraph/src/kagent/langgraph/_executor.py +++ b/python/packages/kagent-langgraph/src/kagent/langgraph/_executor.py @@ -40,9 +40,11 @@ get_tool_approval_response, hitl_activated, now_timestamp, + read_message_metadata, require_ask_user_response, require_tool_approval_response, ) +from kagent.core.tracing import caller_context_attributes from kagent.core.tracing._span_processor import ( clear_kagent_span_attributes, set_kagent_span_attributes, @@ -541,4 +543,8 @@ def _convert_a2a_request_to_span_attributes( if request.task_id: span_attributes["gen_ai.task.id"] = request.task_id + # Allowlisted caller context joins the request-scoped bag rather than a + # single span, so tool, sub-agent, and model spans all carry it. + span_attributes.update(caller_context_attributes(read_message_metadata(request.message))) + return span_attributes From 43abb919f5d1849ab330422a992af25d6fb67a08 Mon Sep 17 00:00:00 2001 From: Ricardo Temperini <29879569+rtemperini@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:31:57 +0200 Subject: [PATCH 03/13] feat(core): forward the trace context allowlist to agent runtimes Expose the allowlist as the Helm value otel.tracing.contextKeys, rendered into the controller ConfigMap as KAGENT_TRACE_CONTEXT_KEYS and forwarded to the agents the controller creates. The variable needs explicit forwarding because collectOtelEnvFromProcess carries only OTEL_ prefixed names. Which caller-supplied data reaches a trace backend is cluster-wide operator policy, so the value is applied after the Harness environment, and any inherited entry of the same name is dropped first. Without that second step a Harness could enable promotion whenever the operator had configured nothing. The value defaults to an empty list, so the ConfigMap key is absent and the runtimes promote nothing unless an operator opts in. Signed-off-by: Ricardo Temperini <29879569+rtemperini@users.noreply.github.com> Co-authored-by: Cursor --- go/core/pkg/env/otel.go | 9 +++ go/core/v2/translator/compiler_test.go | 63 +++++++++++++++++++ go/core/v2/translator/kagent/compiler.go | 13 ++++ .../templates/controller-configmap.yaml | 3 + .../tests/controller-deployment_test.yaml | 18 ++++++ helm/kagent/values.yaml | 5 ++ 6 files changed, 111 insertions(+) diff --git a/go/core/pkg/env/otel.go b/go/core/pkg/env/otel.go index cf3dc1b690..608566a849 100644 --- a/go/core/pkg/env/otel.go +++ b/go/core/pkg/env/otel.go @@ -37,4 +37,13 @@ var ( "OTLP exporter endpoint for logs. Takes precedence over OTEL_EXPORTER_OTLP_ENDPOINT for logs.", ComponentController, ) + + KagentTraceContextKeys = RegisterStringVar( + "KAGENT_TRACE_CONTEXT_KEYS", + "", + "Comma-separated allowlist of caller-supplied context keys promoted onto every agent span as "+ + "kagent.context.. Values are read from W3C baggage and A2A message metadata. "+ + "Empty (the default) disables promotion.", + ComponentAgentRuntime, + ) ) diff --git a/go/core/v2/translator/compiler_test.go b/go/core/v2/translator/compiler_test.go index 771bff0ca0..da28cc872f 100644 --- a/go/core/v2/translator/compiler_test.go +++ b/go/core/v2/translator/compiler_test.go @@ -9,6 +9,7 @@ import ( "github.com/kagent-dev/kagent/go/api/adk" "github.com/kagent-dev/kagent/go/api/v1alpha3" + "github.com/kagent-dev/kagent/go/core/pkg/env" v2translator "github.com/kagent-dev/kagent/go/core/v2/translator" kagenttranslator "github.com/kagent-dev/kagent/go/core/v2/translator/kagent" "github.com/stretchr/testify/require" @@ -262,6 +263,68 @@ func TestCompileAgentTemplateSharedAgent(t *testing.T) { require.Contains(t, string(revision.Provenance), `"name":"researcher"`) } +// KAGENT_TRACE_CONTEXT_KEYS is not an OTEL_ variable, so collectOtelEnvFromProcess +// does not carry it and it needs forwarding of its own. It is also operator +// policy, so a Harness must not be able to widen or enable it. +func TestCompileAgentTemplateForwardsTraceContextKeys(t *testing.T) { + template := &v1alpha3.AgentTemplate{ + ObjectMeta: metav1.ObjectMeta{Name: "helper", Namespace: "test"}, + Spec: v1alpha3.AgentTemplateSpec{ModelConfig: v1alpha3.AgentTemplateLocalReference{Name: "default-model"}}, + } + harnessSupplied := "tenant.supplied" + + tests := []struct { + name string + configured string + harnessEnv []v1alpha3.HarnessEnvVar + want string + }{ + {name: "absent when unconfigured"}, + {name: "forwarded when configured", configured: "user.email,thread_id", want: "user.email,thread_id"}, + { + name: "harness cannot widen the allowlist", + configured: "user.email", + harnessEnv: []v1alpha3.HarnessEnvVar{{Name: env.KagentTraceContextKeys.Name(), Value: &harnessSupplied}}, + want: "user.email", + }, + { + name: "harness cannot enable promotion", + harnessEnv: []v1alpha3.HarnessEnvVar{{Name: env.KagentTraceContextKeys.Name(), Value: &harnessSupplied}}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Setenv(env.KagentTraceContextKeys.Name(), tt.configured) + harness := &v1alpha3.Harness{ + ObjectMeta: metav1.ObjectMeta{Name: "kagent", Namespace: "test"}, + Spec: v1alpha3.HarnessSpec{ + Kagent: &v1alpha3.KagentHarness{}, + AllowedAgentTemplates: &v1alpha3.HarnessAgentTemplateAdmission{Selector: metav1.LabelSelector{}}, + Workload: v1alpha3.HarnessWorkload{Image: "example.com/kagent@sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"}, + Substrate: v1alpha3.HarnessSubstratePolicy{ + WorkerPoolRef: corev1.LocalObjectReference{Name: "default"}, SnapshotPolicy: v1alpha3.HarnessSnapshotPolicy{Location: "snapshots"}, + }, + Env: tt.harnessEnv, + }, + } + + revision, err := compiler(t, modelConfig()).CompileAgentTemplate(context.Background(), harness, template) + require.NoError(t, err) + + var got string + var seen int + for _, variable := range revision.Environment { + if variable.Name == env.KagentTraceContextKeys.Name() { + got, seen = variable.Value, seen+1 + } + } + require.LessOrEqual(t, seen, 1, "environment must not contain a duplicate entry") + require.Equal(t, tt.want, got) + }) + } +} + func TestCompileAgentTemplateRejectsInvalidSharedTrees(t *testing.T) { selector := &v1alpha3.HarnessAgentTemplateAdmission{Selector: metav1.LabelSelector{MatchLabels: map[string]string{"runtime": "kagent"}}} harness := &v1alpha3.Harness{ObjectMeta: metav1.ObjectMeta{Name: "kagent", Namespace: "test"}, Spec: v1alpha3.HarnessSpec{ diff --git a/go/core/v2/translator/kagent/compiler.go b/go/core/v2/translator/kagent/compiler.go index 5be6d58905..2e26147909 100644 --- a/go/core/v2/translator/kagent/compiler.go +++ b/go/core/v2/translator/kagent/compiler.go @@ -91,6 +91,19 @@ func (c *Compiler) Compile(ctx context.Context, input *v2translator.HarnessInput corev1.EnvVar{Name: "KAGENT_A2A_GRPC_ADDRESS", Value: "[::]:80"}, corev1.EnvVar{Name: "KAGENT_PRE_RESPONSE_TRACE_FLUSH", Value: "true"}, ) + // Which caller-supplied context reaches traces is cluster-wide operator + // policy, so a Harness must be able to neither widen nor enable it. Dropping + // any inherited entry before applying the operator's value is what makes that + // hold when the operator has configured nothing at all. + environment = slices.DeleteFunc(environment, func(variable corev1.EnvVar) bool { + return variable.Name == env.KagentTraceContextKeys.Name() + }) + if traceContextKeys := env.KagentTraceContextKeys.Get(); traceContextKeys != "" { + environment = append(environment, corev1.EnvVar{ + Name: env.KagentTraceContextKeys.Name(), + Value: traceContextKeys, + }) + } environment = dedupeEnv(environment) // One provenance list covers every Kubernetes input, including hashed Secret diff --git a/helm/kagent/templates/controller-configmap.yaml b/helm/kagent/templates/controller-configmap.yaml index d31b3e8456..32a66c833c 100644 --- a/helm/kagent/templates/controller-configmap.yaml +++ b/helm/kagent/templates/controller-configmap.yaml @@ -20,6 +20,9 @@ data: # OpenTelemetry Configuration OTEL_TRACING_ENABLED: {{ .Values.otel.tracing.enabled | quote }} OTEL_LOGGING_ENABLED: {{ .Values.otel.logging.enabled | quote }} + {{- with .Values.otel.tracing.contextKeys }} + KAGENT_TRACE_CONTEXT_KEYS: {{ join "," . | quote }} + {{- end }} {{- $tracesEndpoint := .Values.otel.tracing.exporter.otlp.endpoint }} {{- $logsEndpoint := .Values.otel.logging.exporter.otlp.endpoint }} {{- if and $tracesEndpoint $logsEndpoint (eq $tracesEndpoint $logsEndpoint) }} diff --git a/helm/kagent/tests/controller-deployment_test.yaml b/helm/kagent/tests/controller-deployment_test.yaml index a30e1619ce..c94fdf3785 100644 --- a/helm/kagent/tests/controller-deployment_test.yaml +++ b/helm/kagent/tests/controller-deployment_test.yaml @@ -831,3 +831,21 @@ tests: content: name: METRICS_BIND_ADDRESS value: "0" + + - it: should omit KAGENT_TRACE_CONTEXT_KEYS by default + template: controller-configmap.yaml + asserts: + - notExists: + path: data.KAGENT_TRACE_CONTEXT_KEYS + + - it: should join otel.tracing.contextKeys into KAGENT_TRACE_CONTEXT_KEYS + template: controller-configmap.yaml + set: + otel.tracing.contextKeys: + - user.email + - user.name + - thread_id + asserts: + - equal: + path: data.KAGENT_TRACE_CONTEXT_KEYS + value: "user.email,user.name,thread_id" diff --git a/helm/kagent/values.yaml b/helm/kagent/values.yaml index daa3fe9f85..49074317c4 100644 --- a/helm/kagent/values.yaml +++ b/helm/kagent/values.yaml @@ -848,6 +848,11 @@ oauth2-proxy: otel: tracing: enabled: false + # Allowlist of caller-supplied context keys promoted onto every agent span + # as kagent.context.. Values are read from W3C baggage and A2A message + # metadata. Empty (the default) disables promotion. + # e.g. ["user.email", "user.name", "thread_id", "channel"] + contextKeys: [] exporter: otlp: endpoint: "" From efa81704322003833ac6c560eb4d823ca4bb7db6 Mon Sep 17 00:00:00 2001 From: Ricardo Temperini <29879569+rtemperini@users.noreply.github.com> Date: Tue, 25 Aug 2026 18:43:19 +0200 Subject: [PATCH 04/13] docs(architecture): document caller context in traces Cover the configuration knob, why baggage is the primary propagation mechanism, why the attributes are stamped on every span rather than the root, the safety properties that bound untrusted caller input, and how to rename attributes in the OTel Collector for a backend that expects its own names. Signed-off-by: Ricardo Temperini <29879569+rtemperini@users.noreply.github.com> Co-authored-by: Cursor --- docs/architecture/README.md | 1 + docs/architecture/trace-context.md | 136 +++++++++++++++++++++++++++++ 2 files changed, 137 insertions(+) create mode 100644 docs/architecture/trace-context.md diff --git a/docs/architecture/README.md b/docs/architecture/README.md index a0de4fd65b..89a6376c0c 100644 --- a/docs/architecture/README.md +++ b/docs/architecture/README.md @@ -12,6 +12,7 @@ This directory contains detailed architecture documentation for kagent. Start wi | [prompt-templates.md](prompt-templates.md) | Prompt template system with ConfigMap includes and variable interpolation | | [data-flow.md](data-flow.md) | End-to-end request flow from UI to agent and back | | [crds-and-types.md](crds-and-types.md) | All Custom Resource Definitions and their relationships | +| [trace-context.md](trace-context.md) | Promoting caller identity and context onto agent spans | --- diff --git a/docs/architecture/trace-context.md b/docs/architecture/trace-context.md new file mode 100644 index 0000000000..faeb9bae0b --- /dev/null +++ b/docs/architecture/trace-context.md @@ -0,0 +1,136 @@ +# Caller Context in Traces + +Agent spans describe *what the agent did*, but they say nothing about *who asked +for it*. Kagent can promote a configurable allowlist of caller-supplied values — +the signed-in user's email, a Slack thread, a support ticket ID — onto every +span of a request, so traces can be filtered and grouped by the caller in +Langfuse, Jaeger, Grafana Tempo, or any other OTLP backend. + +The feature is **off by default**. It turns on when an operator sets an +allowlist. + +--- + +## Configuration + +| Setting | Default | Description | +|---|---|---| +| Helm `otel.tracing.contextKeys` | `[]` | List of context keys to promote | +| Env `KAGENT_TRACE_CONTEXT_KEYS` | `""` | Comma-separated form of the same list | + +```yaml +otel: + tracing: + enabled: true + contextKeys: + - user.email + - user.name + - thread_id + - channel +``` + +The controller forwards `KAGENT_TRACE_CONTEXT_KEYS` to every agent it creates. +Both the Go and the Python runtime read it, so behaviour is identical whichever +one an agent runs. + +Adding a new traced value is a configuration change, not a code change: append +the key and redeploy. + +--- + +## Where values come from + +Two sources feed the allowlist, in increasing order of precedence: + +| Source | Set by | Survives hops | +|---|---|---| +| W3C [Baggage](https://www.w3.org/TR/baggage/) (`baggage` header) | Any client or proxy on the request path | Yes — automatically | +| A2A `message.metadata` | The A2A caller, per message | No — one hop only | + +**Baggage is the primary mechanism.** It is the vendor-neutral OTel answer to +this problem and it needs no kagent-specific knowledge from the caller: the +controller, both runtimes, and every instrumented HTTP client already run a +composite `tracecontext + baggage` propagator, so a value set once at the edge +reaches the agent, its sub-agents, its tools, and its model calls without any +further plumbing. + +**A2A `message.metadata` is the complement** for callers that cannot set a +header — for example, a bot that speaks A2A over an SDK that exposes message +metadata but not transport headers. It is scoped to a single message, and +because it is the more specific source, it overrides baggage for the same key. + +A key absent from both sources is simply not emitted. + +--- + +## Where values land + +Every promoted value becomes a span attribute named `kagent.context.`: + +``` +baggage: user.email=ada@example.com → kagent.context.user.email = "ada@example.com" +metadata: {"thread_id": "1717171.42"} → kagent.context.thread_id = "1717171.42" +``` + +The attributes are merged into the **request-scoped attribute bag**, not set on +a single span. The `KagentAttributesSpanProcessor` (Python) and +`kagentAttributesSpanProcessor` (Go) stamp that bag onto every span started +during the request, so tool calls, sub-agent delegations, MCP calls, and model +calls all carry the same values. + +This is deliberate rather than incidental: Langfuse v4 and comparable backends +resolve trace-level filters against the attributes present on each span, so +stamping only the root span would leave most views unfilterable. + +--- + +## Safety properties + +Caller-supplied context is untrusted input, so promotion is constrained on every +axis: + +| Risk | Control | +|---|---| +| Attribute explosion / cardinality | Only allowlisted keys are read; the allowlist itself is capped at 32 entries | +| Oversized spans | Values are truncated to 256 characters, keys to 64 | +| Log or trace injection | Control characters are stripped from values | +| Shadowing semantic conventions | Every attribute is namespaced under `kagent.context.`, so `service.name` and friends cannot be overwritten even if an operator allowlists them | +| Leaking secrets into a trace backend | Nothing is promoted unless an operator names the key; raw values are never logged | +| A tenant widening the allowlist | The allowlist is cluster-wide operator configuration; an entry of the same name in a `Harness` environment is dropped rather than inherited | + +Non-scalar metadata (objects, arrays) is skipped: it is unbounded in size and +meaningless as an attribute value. + +Choose allowlist keys deliberately. Anything named here is visible to everyone +with access to the trace backend, and callers control the values. + +--- + +## Renaming attributes for a backend + +Some backends expect specific attribute names — Langfuse, for instance, maps +`user.id` and `session.id` onto its own trace fields. Rather than making the +attribute namespace configurable, do the rename in the OTel Collector that +already sits between kagent and the backend: + +```yaml +processors: + transform: + trace_statements: + - set(span.attributes["user.id"], span.attributes["kagent.context.user.email"]) + where span.attributes["kagent.context.user.email"] != nil +``` + +--- + +## Implementation + +| Component | Path | +|---|---| +| Go ADK | `go/adk/pkg/telemetry/context_attributes.go` | +| Python | `python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py` | +| Controller forwarding | `go/core/v2/translator/kagent/compiler.go` | +| Helm | `helm/kagent/templates/controller-configmap.yaml` | + +Both implementations share the same allowlist parsing, precedence, limits, and +sanitisation rules so the two runtimes cannot drift. From 20565918c990d357b4b3ab4c5aa3229218655e17 Mon Sep 17 00:00:00 2001 From: Ricardo Temperini <29879569+rtemperini@users.noreply.github.com> Date: Fri, 28 Aug 2026 12:20:47 +0200 Subject: [PATCH 05/13] fix: emit registry attributes unprefixed and hash identifiers Address review of the caller-context tracing allowlist: use an OIDC subject as the user.id example, leave user.*, enduser.*, and session.id unprefixed, and hash mapped values with HMAC-SHA256 instead of putting the original on the span. Signed-off-by: Ricardo Temperini <29879569+rtemperini@users.noreply.github.com> Co-authored-by: Cursor --- docs/architecture/trace-context.md | 75 +++++-- go/adk/pkg/telemetry/context_attributes.go | 206 ++++++++++++++---- .../pkg/telemetry/context_attributes_test.go | 160 +++++++++++--- go/core/pkg/env/otel.go | 16 +- go/core/v2/translator/compiler_test.go | 58 +++-- go/core/v2/translator/kagent/compiler.go | 25 ++- helm/kagent/templates/_helpers.tpl | 17 ++ .../templates/controller-configmap.yaml | 2 +- .../templates/controller-deployment.yaml | 9 + .../tests/controller-deployment_test.yaml | 44 +++- helm/kagent/values.yaml | 28 ++- .../core/tracing/_context_attributes.py | 180 ++++++++++++--- .../tests/test_caller_context_attributes.py | 122 +++++++++-- .../tests/test_tracing_configure.py | 4 +- 14 files changed, 772 insertions(+), 174 deletions(-) diff --git a/docs/architecture/trace-context.md b/docs/architecture/trace-context.md index faeb9bae0b..aed95a69a8 100644 --- a/docs/architecture/trace-context.md +++ b/docs/architecture/trace-context.md @@ -2,9 +2,9 @@ Agent spans describe *what the agent did*, but they say nothing about *who asked for it*. Kagent can promote a configurable allowlist of caller-supplied values — -the signed-in user's email, a Slack thread, a support ticket ID — onto every -span of a request, so traces can be filtered and grouped by the caller in -Langfuse, Jaeger, Grafana Tempo, or any other OTLP backend. +an opaque user identifier, a conversation thread, a ticket ID — onto every span +of a request, so traces can be filtered and grouped by the caller in Langfuse, +Jaeger, Grafana Tempo, or any other OTLP backend. The feature is **off by default**. It turns on when an operator sets an allowlist. @@ -15,20 +15,23 @@ allowlist. | Setting | Default | Description | |---|---|---| -| Helm `otel.tracing.contextKeys` | `[]` | List of context keys to promote | -| Env `KAGENT_TRACE_CONTEXT_KEYS` | `""` | Comma-separated form of the same list | +| Helm `otel.tracing.contextKeys` | `[]` | List of context keys or `{from, to, hash}` mappings to promote | +| Env `KAGENT_TRACE_CONTEXT_KEYS` | `""` | Comma-separated keys, or a JSON array of the same mappings | +| Helm `otel.tracing.contextHashKeySecret` | unset | Secret providing `KAGENT_TRACE_CONTEXT_HASH_KEY` for `hash: hmac-sha256` | ```yaml otel: tracing: enabled: true contextKeys: - - user.email - - user.name - - thread_id + - {from: sub, to: user.id} + - {from: thread_id, to: kagent.thread_id} - channel ``` +Prefer an opaque identifier such as an OIDC `sub` for `user.id`. Do not put +names or email addresses on spans; see [Sensitive values](#sensitive-values). + The controller forwards `KAGENT_TRACE_CONTEXT_KEYS` to every agent it creates. Both the Go and the Python runtime read it, so behaviour is identical whichever one an agent runs. @@ -65,13 +68,21 @@ A key absent from both sources is simply not emitted. ## Where values land -Every promoted value becomes a span attribute named `kagent.context.`: +Each mapping is read from `from` (defaulting to the entry itself) and written +as span attribute `to` (defaulting to `from`) after the prefix rules below: ``` -baggage: user.email=ada@example.com → kagent.context.user.email = "ada@example.com" -metadata: {"thread_id": "1717171.42"} → kagent.context.thread_id = "1717171.42" +baggage: sub=opaque-subject → user.id = "opaque-subject" +metadata: {"thread_id": "1717171.42"} → kagent.thread_id = "1717171.42" +metadata: {"channel": "C0AB1"} → kagent.context.channel = "C0AB1" ``` +| Destination name | Emitted as | +|---|---| +| `user.*`, `enduser.*`, `session.id` | Unprefixed (OpenTelemetry semantic conventions) | +| Already in the `kagent.` namespace | Unprefixed | +| Anything else | `kagent.context.` | + The attributes are merged into the **request-scoped attribute bag**, not set on a single span. The `KagentAttributesSpanProcessor` (Python) and `kagentAttributesSpanProcessor` (Go) stamp that bag onto every span started @@ -84,6 +95,35 @@ stamping only the root span would leave most views unfilterable. --- +## Sensitive values + +[OpenTelemetry recommends](https://opentelemetry.io/docs/security/handling-sensitive-data/) +against putting email addresses or names on telemetry at all. An OIDC `sub` is +already an opaque identifier and is what `user.id` should use. + +If a stable identifier must be derived from a value that itself should not +appear on a span, hash it with HMAC-SHA256 onto the registry attribute +`user.hash`: + +```yaml +contextKeys: + - {from: sub, to: user.id} + - {from: email, to: user.hash, hash: hmac-sha256} + - {from: thread_id, to: kagent.thread_id} +``` + +`hash: hmac-sha256` requires `KAGENT_TRACE_CONTEXT_HASH_KEY` (Helm: +`otel.tracing.contextHashKeySecret`). If the key is missing, the hashed +attribute is skipped — the original value is never written onto the span. + +Hashing at promotion time only affects the span. Baggage travels on HTTP +headers, so a value placed in baggage is still visible to every downstream hop +that receives those headers, including model providers and HTTP MCP servers. +Do not put sensitive values in baggage; hash or replace them at the edge +before the request enters the cluster. + +--- + ## Safety properties Caller-supplied context is untrusted input, so promotion is constrained on every @@ -94,8 +134,8 @@ axis: | Attribute explosion / cardinality | Only allowlisted keys are read; the allowlist itself is capped at 32 entries | | Oversized spans | Values are truncated to 256 characters, keys to 64 | | Log or trace injection | Control characters are stripped from values | -| Shadowing semantic conventions | Every attribute is namespaced under `kagent.context.`, so `service.name` and friends cannot be overwritten even if an operator allowlists them | -| Leaking secrets into a trace backend | Nothing is promoted unless an operator names the key; raw values are never logged | +| Shadowing semantic conventions | Custom keys are namespaced under `kagent.context.`; only `user.*`, `enduser.*`, and `session.id` pass through unprefixed | +| Leaking secrets into a trace backend | Nothing is promoted unless an operator names the key; hashed entries are omitted when the HMAC key is unset | | A tenant widening the allowlist | The allowlist is cluster-wide operator configuration; an entry of the same name in a `Harness` environment is dropped rather than inherited | Non-scalar metadata (objects, arrays) is skipped: it is unbounded in size and @@ -108,17 +148,16 @@ with access to the trace backend, and callers control the values. ## Renaming attributes for a backend -Some backends expect specific attribute names — Langfuse, for instance, maps -`user.id` and `session.id` onto its own trace fields. Rather than making the -attribute namespace configurable, do the rename in the OTel Collector that +Some backends expect names other than the ones kagent emits. Rather than making +the attribute namespace configurable, do the rename in the OTel Collector that already sits between kagent and the backend: ```yaml processors: transform: trace_statements: - - set(span.attributes["user.id"], span.attributes["kagent.context.user.email"]) - where span.attributes["kagent.context.user.email"] != nil + - set(span.attributes["session.id"], span.attributes["kagent.thread_id"]) + where span.attributes["kagent.thread_id"] != nil ``` --- diff --git a/go/adk/pkg/telemetry/context_attributes.go b/go/adk/pkg/telemetry/context_attributes.go index f0de636101..d0774d144f 100644 --- a/go/adk/pkg/telemetry/context_attributes.go +++ b/go/adk/pkg/telemetry/context_attributes.go @@ -2,6 +2,10 @@ package telemetry import ( "context" + "crypto/hmac" + "crypto/sha256" + "encoding/hex" + "encoding/json" "os" "strconv" "strings" @@ -11,21 +15,40 @@ import ( ) const ( - // traceContextKeysEnvVar holds a comma-separated allowlist of caller-supplied - // context keys to promote onto agent spans. Unset or empty (the default) - // disables promotion entirely. + // traceContextKeysEnvVar holds the allowlist of caller-supplied context + // keys to promote onto agent spans. Unset or empty (the default) disables + // promotion entirely. + // + // Accepts a comma-separated list of source keys, or a JSON array of strings + // and {from, to, hash} objects. See allowedContextMappings. traceContextKeysEnvVar = "KAGENT_TRACE_CONTEXT_KEYS" - // contextAttributePrefix namespaces every promoted value. Because the prefix - // is applied unconditionally, caller-supplied data cannot shadow a semantic - // convention attribute such as service.name. + // traceContextHashKeyEnvVar is the HMAC key used when a mapping sets + // hash: hmac-sha256. Required for those entries; without it the hashed + // attribute is skipped rather than emitted in plaintext. + traceContextHashKeyEnvVar = "KAGENT_TRACE_CONTEXT_HASH_KEY" + + // contextAttributePrefix namespaces custom promoted values so they cannot + // shadow a semantic convention attribute such as service.name. Registry + // names (user.*, enduser.*, session.id) and names already in the kagent. + // namespace are left unprefixed; see spanAttributeName. contextAttributePrefix = "kagent.context." + hashHMACSHA256 = "hmac-sha256" + maxContextKeys = 32 maxContextKeyLength = 64 maxContextValueLength = 256 ) +// contextMapping is one allowlisted promotion: read source from baggage or +// A2A metadata and emit it as span attribute attribute (after prefix rules). +type contextMapping struct { + source string + attribute string + hash string +} + // CallerContextAttributes returns the caller-supplied context values that an // operator allowlisted through KAGENT_TRACE_CONTEXT_KEYS. Merge the result into // the request-scoped attribute bag (see SetKAgentSpanAttributes) so every span @@ -35,27 +58,39 @@ const ( // Values are read from W3C baggage first and then from the A2A message // metadata, which is the more specific source for a single message and // therefore wins. Both are untrusted input, so keys must appear in the -// allowlist, values are stripped of control characters and truncated, and every -// attribute is namespaced under contextAttributePrefix. +// allowlist, values are stripped of control characters and truncated, and +// attribute names go through spanAttributeName. // // Returns nil when the allowlist is empty, which is the default. func CallerContextAttributes(ctx context.Context, metadata map[string]any) map[string]string { - keys := allowedContextKeys() - if len(keys) == 0 { + mappings := allowedContextMappings() + if len(mappings) == 0 { return nil } bag := baggage.FromContext(ctx) - attrs := make(map[string]string, len(keys)) - for _, key := range keys { - value := sanitizeContextValue(bag.Member(key).Value()) - if scalar, ok := scalarString(metadata[key]); ok { + attrs := make(map[string]string, len(mappings)) + for _, mapping := range mappings { + value := sanitizeContextValue(bag.Member(mapping.source).Value()) + if scalar, ok := scalarString(metadata[mapping.source]); ok { value = sanitizeContextValue(scalar) } if value == "" { continue } - attrs[contextAttributePrefix+key] = value + if mapping.hash != "" { + value = hashContextValue(value, mapping.hash) + if value == "" { + continue + } + } else { + value = truncateContextValue(value) + } + name := spanAttributeName(mapping.attribute) + if _, exists := attrs[name]; exists { + continue + } + attrs[name] = value } if len(attrs) == 0 { return nil @@ -63,33 +98,113 @@ func CallerContextAttributes(ctx context.Context, metadata map[string]any) map[s return attrs } -// allowedContextKeys parses the KAGENT_TRACE_CONTEXT_KEYS allowlist. Keys that -// are empty, over-long, or contain whitespace or control characters are -// dropped, and the list is capped at maxContextKeys so a misconfigured -// allowlist cannot inflate span cardinality without bound. -func allowedContextKeys() []string { +// allowedContextMappings parses the KAGENT_TRACE_CONTEXT_KEYS allowlist. +// Entries that are empty, over-long, or contain whitespace or control +// characters are dropped, and the list is capped at maxContextKeys so a +// misconfigured allowlist cannot inflate span cardinality without bound. +func allowedContextMappings() []contextMapping { raw := strings.TrimSpace(os.Getenv(traceContextKeysEnvVar)) if raw == "" { return nil } + if strings.HasPrefix(raw, "[") { + return capMappings(parseJSONAllowlist(raw)) + } + return capMappings(parseCommaAllowlist(raw)) +} - keys := make([]string, 0, maxContextKeys) - seen := make(map[string]struct{}, maxContextKeys) +func parseCommaAllowlist(raw string) []contextMapping { + mappings := make([]contextMapping, 0, maxContextKeys) for key := range strings.SplitSeq(raw, ",") { - key = strings.TrimSpace(key) - if key == "" || len(key) > maxContextKeyLength || !isAttributeKey(key) { + if mapping, ok := newContextMapping(strings.TrimSpace(key), "", ""); ok { + mappings = append(mappings, mapping) + } + } + return mappings +} + +func parseJSONAllowlist(raw string) []contextMapping { + var items []json.RawMessage + if err := json.Unmarshal([]byte(raw), &items); err != nil { + return nil + } + mappings := make([]contextMapping, 0, maxContextKeys) + for _, item := range items { + var key string + if err := json.Unmarshal(item, &key); err == nil { + if mapping, ok := newContextMapping(key, "", ""); ok { + mappings = append(mappings, mapping) + } continue } - if _, duplicate := seen[key]; duplicate { + var spec struct { + From string `json:"from"` + To string `json:"to"` + Hash string `json:"hash"` + } + if err := json.Unmarshal(item, &spec); err != nil { + continue + } + if mapping, ok := newContextMapping(spec.From, spec.To, spec.Hash); ok { + mappings = append(mappings, mapping) + } + } + return mappings +} + +func newContextMapping(from, to, hash string) (contextMapping, bool) { + from = strings.TrimSpace(from) + to = strings.TrimSpace(to) + hash = strings.TrimSpace(hash) + if from == "" || len(from) > maxContextKeyLength || !isAttributeKey(from) { + return contextMapping{}, false + } + if to == "" { + to = from + } + if len(to) > maxContextKeyLength || !isAttributeKey(to) { + return contextMapping{}, false + } + if hash != "" && hash != hashHMACSHA256 { + return contextMapping{}, false + } + return contextMapping{source: from, attribute: to, hash: hash}, true +} + +func capMappings(mappings []contextMapping) []contextMapping { + out := make([]contextMapping, 0, maxContextKeys) + seen := make(map[string]struct{}, maxContextKeys) + for _, mapping := range mappings { + id := mapping.source + "\x00" + mapping.attribute + "\x00" + mapping.hash + if _, duplicate := seen[id]; duplicate { continue } - seen[key] = struct{}{} - keys = append(keys, key) - if len(keys) == maxContextKeys { + seen[id] = struct{}{} + out = append(out, mapping) + if len(out) == maxContextKeys { break } } - return keys + return out +} + +// spanAttributeName is the name written onto the span. +// +// user.*, enduser.*, and session.id pass through unprefixed so operators can +// use the semantic convention names. Names already in the kagent. namespace +// are left as-is. Everything else is placed under kagent.context. so a +// caller-supplied service.name cannot shadow the real one. +func spanAttributeName(name string) string { + if isRegistryAttribute(name) || strings.HasPrefix(name, "kagent.") { + return name + } + return contextAttributePrefix + name +} + +func isRegistryAttribute(name string) bool { + return strings.HasPrefix(name, "user.") || + strings.HasPrefix(name, "enduser.") || + name == "session.id" } // isAttributeKey reports whether key is safe to use as a span attribute name. @@ -119,10 +234,8 @@ func scalarString(value any) (string, bool) { } } -// sanitizeContextValue makes an untrusted value safe to attach to a span: -// control characters are dropped so a value cannot forge structure in a -// downstream trace or log renderer, and the result is truncated to bound the -// size of exported spans. +// sanitizeContextValue drops control characters and trims space so a value +// cannot forge structure in a downstream trace or log renderer. func sanitizeContextValue(value string) string { cleaned := strings.Map(func(r rune) rune { if unicode.IsControl(r) { @@ -130,11 +243,30 @@ func sanitizeContextValue(value string) string { } return r }, value) - cleaned = strings.TrimSpace(cleaned) + return strings.TrimSpace(cleaned) +} + +func truncateContextValue(value string) string { // Truncate by rune, not byte, so the limit means the same thing here as it // does in the Python runtime. - if runes := []rune(cleaned); len(runes) > maxContextValueLength { - cleaned = string(runes[:maxContextValueLength]) + if runes := []rune(value); len(runes) > maxContextValueLength { + return string(runes[:maxContextValueLength]) + } + return value +} + +// hashContextValue hashes value with the requested algorithm. Unknown +// algorithms and a missing HMAC key skip the attribute: never fall back to +// putting the original value on the span. +func hashContextValue(value, algorithm string) string { + if algorithm != hashHMACSHA256 { + return "" + } + key := os.Getenv(traceContextHashKeyEnvVar) + if key == "" { + return "" } - return cleaned + mac := hmac.New(sha256.New, []byte(key)) + _, _ = mac.Write([]byte(value)) + return hex.EncodeToString(mac.Sum(nil)) } diff --git a/go/adk/pkg/telemetry/context_attributes_test.go b/go/adk/pkg/telemetry/context_attributes_test.go index 5a4eb06c6c..89d9962184 100644 --- a/go/adk/pkg/telemetry/context_attributes_test.go +++ b/go/adk/pkg/telemetry/context_attributes_test.go @@ -2,6 +2,9 @@ package telemetry import ( "context" + "crypto/hmac" + "crypto/sha256" + "encoding/hex" "strconv" "strings" "testing" @@ -29,28 +32,36 @@ func baggageContext(t *testing.T, members map[string]string) context.Context { return baggage.ContextWithBaggage(context.Background(), bag) } +func hmacSHA256Hex(t *testing.T, key, value string) string { + t.Helper() + mac := hmac.New(sha256.New, []byte(key)) + _, _ = mac.Write([]byte(value)) + return hex.EncodeToString(mac.Sum(nil)) +} + func TestCallerContextAttributes(t *testing.T) { tests := []struct { name string allowlist string + hashKey string baggageVals map[string]string metadata map[string]any want map[string]string }{ { - name: "disabled by default", + name: "empty allowlist disables promotion", allowlist: "", - baggageVals: map[string]string{"user.email": "ada@example.com"}, + baggageVals: map[string]string{"sub": "opaque-subject"}, metadata: map[string]any{"thread_id": "T123"}, want: nil, }, { name: "promotes allowlisted baggage", - allowlist: "user.email,user.name", - baggageVals: map[string]string{"user.email": "ada@example.com", "user.name": "Ada"}, + allowlist: "sub,thread_id", + baggageVals: map[string]string{"sub": "opaque-subject", "thread_id": "T123"}, want: map[string]string{ - "kagent.context.user.email": "ada@example.com", - "kagent.context.user.name": "Ada", + "kagent.context.sub": "opaque-subject", + "kagent.context.thread_id": "T123", }, }, { @@ -64,16 +75,16 @@ func TestCallerContextAttributes(t *testing.T) { }, { name: "message metadata overrides baggage", - allowlist: "user.email", - baggageVals: map[string]string{"user.email": "from-baggage@example.com"}, - metadata: map[string]any{"user.email": "from-metadata@example.com"}, - want: map[string]string{"kagent.context.user.email": "from-metadata@example.com"}, + allowlist: "sub", + baggageVals: map[string]string{"sub": "from-baggage"}, + metadata: map[string]any{"sub": "from-metadata"}, + want: map[string]string{"kagent.context.sub": "from-metadata"}, }, { name: "ignores keys outside the allowlist", allowlist: "thread_id", baggageVals: map[string]string{"secret.token": "s3cret"}, - metadata: map[string]any{"thread_id": "T1", "customer.pan": "4111111111111111"}, + metadata: map[string]any{"thread_id": "T1", "extra": "nope"}, want: map[string]string{"kagent.context.thread_id": "T1"}, }, { @@ -112,11 +123,62 @@ func TestCallerContextAttributes(t *testing.T) { metadata: map[string]any{"good": "yes", "bad key": "no"}, want: map[string]string{"kagent.context.good": "yes"}, }, + { + // Registry names pass through unprefixed so operators can use + // semantic convention attributes instead of inventing new ones. + name: "registry attributes stay unprefixed", + allowlist: "user.id,enduser.id,session.id,channel", + metadata: map[string]any{ + "user.id": "opaque-subject", + "enduser.id": "end-user", + "session.id": "sess-1", + "channel": "C0AB1", + }, + want: map[string]string{ + "user.id": "opaque-subject", + "enduser.id": "end-user", + "session.id": "sess-1", + "kagent.context.channel": "C0AB1", + }, + }, + { + // session.id is the registry name; other session.* keys are not. + name: "session.id is unprefixed but session.foo is not", + allowlist: "session.id,session.foo", + metadata: map[string]any{"session.id": "sess-1", "session.foo": "other"}, + want: map[string]string{ + "session.id": "sess-1", + "kagent.context.session.foo": "other", + }, + }, + { + name: "maps source keys onto registry and kagent names", + allowlist: `[{"from":"sub","to":"user.id"},{"from":"thread_id","to":"kagent.thread_id"},"channel"]`, + metadata: map[string]any{ + "sub": "opaque-subject", + "thread_id": "T123", + "channel": "C0AB1", + }, + want: map[string]string{ + "user.id": "opaque-subject", + "kagent.thread_id": "T123", + "kagent.context.channel": "C0AB1", + }, + }, + { + name: "invalid JSON allowlist promotes nothing", + allowlist: `[{"from":"sub"`, + metadata: map[string]any{"sub": "opaque-subject"}, + want: nil, + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Setenv(traceContextKeysEnvVar, tt.allowlist) + if tt.hashKey != "" { + t.Setenv(traceContextHashKeyEnvVar, tt.hashKey) + } got := CallerContextAttributes(baggageContext(t, tt.baggageVals), tt.metadata) @@ -132,6 +194,51 @@ func TestCallerContextAttributes(t *testing.T) { } } +func TestCallerContextAttributes_HashesWithHMACSHA256(t *testing.T) { + const key = "test-hmac-key" + t.Setenv(traceContextKeysEnvVar, `[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]`) + t.Setenv(traceContextHashKeyEnvVar, key) + + got := CallerContextAttributes(context.Background(), map[string]any{ + "email": "ada@example.com", + }) + + want := hmacSHA256Hex(t, key, "ada@example.com") + if got["user.hash"] != want { + t.Errorf("user.hash = %q, want %q", got["user.hash"], want) + } + for name, value := range got { + if strings.Contains(value, "@example.com") { + t.Errorf("plaintext leaked onto %s", name) + } + } +} + +func TestCallerContextAttributes_HashWithoutKeyEmitsNothing(t *testing.T) { + // Missing HMAC key must not fall back to putting the original value on the span. + t.Setenv(traceContextKeysEnvVar, `[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]`) + t.Setenv(traceContextHashKeyEnvVar, "") + + got := CallerContextAttributes(context.Background(), map[string]any{ + "email": "ada@example.com", + }) + if got != nil { + t.Errorf("got %v, want nothing when the HMAC key is unset", got) + } +} + +func TestCallerContextAttributes_UnknownHashEmitsNothing(t *testing.T) { + t.Setenv(traceContextKeysEnvVar, `[{"from":"email","to":"user.hash","hash":"md5"}]`) + t.Setenv(traceContextHashKeyEnvVar, "test-hmac-key") + + got := CallerContextAttributes(context.Background(), map[string]any{ + "email": "ada@example.com", + }) + if got != nil { + t.Errorf("got %v, want nothing for an unsupported hash", got) + } +} + func TestCallerContextAttributes_TruncatesLongValues(t *testing.T) { t.Setenv(traceContextKeysEnvVar, "note") @@ -144,30 +251,30 @@ func TestCallerContextAttributes_TruncatesLongValues(t *testing.T) { } } -func TestAllowedContextKeys_CapsListLength(t *testing.T) { +func TestAllowedContextMappings_CapsListLength(t *testing.T) { keys := make([]string, 0, maxContextKeys*2) for i := range maxContextKeys * 2 { keys = append(keys, "key"+strconv.Itoa(i)) } t.Setenv(traceContextKeysEnvVar, strings.Join(keys, ",")) - if got := len(allowedContextKeys()); got != maxContextKeys { + if got := len(allowedContextMappings()); got != maxContextKeys { t.Errorf("allowlist length = %d, want %d", got, maxContextKeys) } } -func TestAllowedContextKeys_DropsOverLongAndDuplicateKeys(t *testing.T) { +func TestAllowedContextMappings_DropsOverLongAndDuplicateKeys(t *testing.T) { t.Setenv(traceContextKeysEnvVar, "a,a,"+strings.Repeat("b", maxContextKeyLength+1)+",c") - got := allowedContextKeys() + got := allowedContextMappings() want := []string{"a", "c"} if len(got) != len(want) { - t.Fatalf("got %v, want %v", got, want) + t.Fatalf("got %#v, want %v", got, want) } for i, key := range want { - if got[i] != key { - t.Errorf("key %d = %q, want %q", i, got[i], key) + if got[i].source != key { + t.Errorf("key %d = %q, want %q", i, got[i].source, key) } } } @@ -175,7 +282,7 @@ func TestAllowedContextKeys_DropsOverLongAndDuplicateKeys(t *testing.T) { // Langfuse and comparable backends filter on attributes present on each span, // so the promoted values must reach descendants, not just the root span. func TestCallerContextAttributes_ReachEverySpan(t *testing.T) { - t.Setenv(traceContextKeysEnvVar, "user.email") + t.Setenv(traceContextKeysEnvVar, `[{"from":"sub","to":"user.id"}]`) exporter := tracetest.NewInMemoryExporter() tp := sdktrace.NewTracerProvider( @@ -186,7 +293,7 @@ func TestCallerContextAttributes_ReachEverySpan(t *testing.T) { _ = tp.Shutdown(context.Background()) }) - ctx := baggageContext(t, map[string]string{"user.email": "ada@example.com"}) + ctx := baggageContext(t, map[string]string{"sub": "opaque-subject"}) ctx = SetKAgentSpanAttributes(ctx, CallerContextAttributes(ctx, nil)) tracer := tp.Tracer("test") @@ -203,8 +310,8 @@ func TestCallerContextAttributes_ReachEverySpan(t *testing.T) { } for _, name := range []string{"root", "execute_tool", "generate_content"} { attrs := spanAttributesByName(t, spans, name) - if got := attrs["kagent.context.user.email"].AsString(); got != "ada@example.com" { - t.Errorf("span %q: kagent.context.user.email = %q, want %q", name, got, "ada@example.com") + if got := attrs["user.id"].AsString(); got != "opaque-subject" { + t.Errorf("span %q: user.id = %q, want %q", name, got, "opaque-subject") } } } @@ -221,21 +328,22 @@ func TestCallerContextAttributes_DisabledLeavesSpansUnchanged(t *testing.T) { _ = tp.Shutdown(context.Background()) }) - ctx := baggageContext(t, map[string]string{"user.email": "ada@example.com"}) + ctx := baggageContext(t, map[string]string{"sub": "opaque-subject"}) ctx = SetKAgentSpanAttributes(ctx, CallerContextAttributes(ctx, map[string]any{"thread_id": "T1"})) _, span := tp.Tracer("test").Start(ctx, "root") span.End() for _, attr := range exporter.GetSpans()[0].Attributes { - if strings.HasPrefix(string(attr.Key), contextAttributePrefix) { + key := string(attr.Key) + if strings.HasPrefix(key, contextAttributePrefix) || key == "user.id" { t.Errorf("unexpected promoted attribute %q", attr.Key) } } } -// The kagent.context. prefix is what makes caller-supplied data unable to -// shadow a semantic convention attribute, even if an operator allowlists one. +// Custom keys still cannot shadow a semantic convention attribute such as +// service.name. Registry names are the documented exception. func TestCallerContextAttributes_CannotShadowSemanticConventions(t *testing.T) { t.Setenv(traceContextKeysEnvVar, "service.name") diff --git a/go/core/pkg/env/otel.go b/go/core/pkg/env/otel.go index 608566a849..2b8775b653 100644 --- a/go/core/pkg/env/otel.go +++ b/go/core/pkg/env/otel.go @@ -41,9 +41,19 @@ var ( KagentTraceContextKeys = RegisterStringVar( "KAGENT_TRACE_CONTEXT_KEYS", "", - "Comma-separated allowlist of caller-supplied context keys promoted onto every agent span as "+ - "kagent.context.. Values are read from W3C baggage and A2A message metadata. "+ - "Empty (the default) disables promotion.", + "Allowlist of caller-supplied context keys promoted onto every agent span. "+ + "Accepts a comma-separated list of source keys, or a JSON array of strings and "+ + "{from, to, hash} objects. Registry names (user.*, enduser.*, session.id) are left "+ + "unprefixed; everything else is emitted as kagent.context. unless the name is "+ + "already in the kagent. namespace. Empty (the default) disables promotion.", + ComponentAgentRuntime, + ) + + KagentTraceContextHashKey = RegisterStringVar( + "KAGENT_TRACE_CONTEXT_HASH_KEY", + "", + "HMAC-SHA256 key used when a KAGENT_TRACE_CONTEXT_KEYS mapping sets hash: hmac-sha256. "+ + "Hashed attributes are omitted when this is unset, rather than emitting the original value.", ComponentAgentRuntime, ) ) diff --git a/go/core/v2/translator/compiler_test.go b/go/core/v2/translator/compiler_test.go index da28cc872f..e23edcbc45 100644 --- a/go/core/v2/translator/compiler_test.go +++ b/go/core/v2/translator/compiler_test.go @@ -263,10 +263,11 @@ func TestCompileAgentTemplateSharedAgent(t *testing.T) { require.Contains(t, string(revision.Provenance), `"name":"researcher"`) } -// KAGENT_TRACE_CONTEXT_KEYS is not an OTEL_ variable, so collectOtelEnvFromProcess -// does not carry it and it needs forwarding of its own. It is also operator -// policy, so a Harness must not be able to widen or enable it. -func TestCompileAgentTemplateForwardsTraceContextKeys(t *testing.T) { +// KAGENT_TRACE_CONTEXT_KEYS and KAGENT_TRACE_CONTEXT_HASH_KEY are not OTEL_ +// variables, so collectOtelEnvFromProcess does not carry them and they need +// forwarding of their own. They are also operator policy, so a Harness must +// not be able to widen, enable, or supply the HMAC key. +func TestCompileAgentTemplateForwardsTraceContextPolicy(t *testing.T) { template := &v1alpha3.AgentTemplate{ ObjectMeta: metav1.ObjectMeta{Name: "helper", Namespace: "test"}, Spec: v1alpha3.AgentTemplateSpec{ModelConfig: v1alpha3.AgentTemplateLocalReference{Name: "default-model"}}, @@ -275,27 +276,47 @@ func TestCompileAgentTemplateForwardsTraceContextKeys(t *testing.T) { tests := []struct { name string - configured string + keys string + hashKey string harnessEnv []v1alpha3.HarnessEnvVar - want string + wantKeys string + wantHash string }{ {name: "absent when unconfigured"}, - {name: "forwarded when configured", configured: "user.email,thread_id", want: "user.email,thread_id"}, + {name: "keys forwarded when configured", keys: "sub,thread_id", wantKeys: "sub,thread_id"}, + { + name: "hash key forwarded when configured", + keys: `[{"from":"sub","to":"user.id"}]`, + hashKey: "test-hmac-key", + wantKeys: `[{"from":"sub","to":"user.id"}]`, + wantHash: "test-hmac-key", + }, { name: "harness cannot widen the allowlist", - configured: "user.email", + keys: "sub", harnessEnv: []v1alpha3.HarnessEnvVar{{Name: env.KagentTraceContextKeys.Name(), Value: &harnessSupplied}}, - want: "user.email", + wantKeys: "sub", }, { name: "harness cannot enable promotion", harnessEnv: []v1alpha3.HarnessEnvVar{{Name: env.KagentTraceContextKeys.Name(), Value: &harnessSupplied}}, }, + { + name: "harness cannot supply the HMAC key", + harnessEnv: []v1alpha3.HarnessEnvVar{{Name: env.KagentTraceContextHashKey.Name(), Value: &harnessSupplied}}, + }, + { + name: "harness cannot replace the HMAC key", + hashKey: "operator-hmac-key", + harnessEnv: []v1alpha3.HarnessEnvVar{{Name: env.KagentTraceContextHashKey.Name(), Value: &harnessSupplied}}, + wantHash: "operator-hmac-key", + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - t.Setenv(env.KagentTraceContextKeys.Name(), tt.configured) + t.Setenv(env.KagentTraceContextKeys.Name(), tt.keys) + t.Setenv(env.KagentTraceContextHashKey.Name(), tt.hashKey) harness := &v1alpha3.Harness{ ObjectMeta: metav1.ObjectMeta{Name: "kagent", Namespace: "test"}, Spec: v1alpha3.HarnessSpec{ @@ -312,15 +333,20 @@ func TestCompileAgentTemplateForwardsTraceContextKeys(t *testing.T) { revision, err := compiler(t, modelConfig()).CompileAgentTemplate(context.Background(), harness, template) require.NoError(t, err) - var got string - var seen int + gotKeys, seenKeys := "", 0 + gotHash, seenHash := "", 0 for _, variable := range revision.Environment { - if variable.Name == env.KagentTraceContextKeys.Name() { - got, seen = variable.Value, seen+1 + switch variable.Name { + case env.KagentTraceContextKeys.Name(): + gotKeys, seenKeys = variable.Value, seenKeys+1 + case env.KagentTraceContextHashKey.Name(): + gotHash, seenHash = variable.Value, seenHash+1 } } - require.LessOrEqual(t, seen, 1, "environment must not contain a duplicate entry") - require.Equal(t, tt.want, got) + require.LessOrEqual(t, seenKeys, 1, "environment must not contain a duplicate keys entry") + require.LessOrEqual(t, seenHash, 1, "environment must not contain a duplicate hash key entry") + require.Equal(t, tt.wantKeys, gotKeys) + require.Equal(t, tt.wantHash, gotHash) }) } } diff --git a/go/core/v2/translator/kagent/compiler.go b/go/core/v2/translator/kagent/compiler.go index 2e26147909..7373b7caec 100644 --- a/go/core/v2/translator/kagent/compiler.go +++ b/go/core/v2/translator/kagent/compiler.go @@ -94,16 +94,10 @@ func (c *Compiler) Compile(ctx context.Context, input *v2translator.HarnessInput // Which caller-supplied context reaches traces is cluster-wide operator // policy, so a Harness must be able to neither widen nor enable it. Dropping // any inherited entry before applying the operator's value is what makes that - // hold when the operator has configured nothing at all. - environment = slices.DeleteFunc(environment, func(variable corev1.EnvVar) bool { - return variable.Name == env.KagentTraceContextKeys.Name() - }) - if traceContextKeys := env.KagentTraceContextKeys.Get(); traceContextKeys != "" { - environment = append(environment, corev1.EnvVar{ - Name: env.KagentTraceContextKeys.Name(), - Value: traceContextKeys, - }) - } + // hold when the operator has configured nothing at all. The HMAC key is the + // same class of policy: a tenant must not supply it. + environment = applyOperatorOnlyEnv(environment, env.KagentTraceContextKeys) + environment = applyOperatorOnlyEnv(environment, env.KagentTraceContextHashKey) environment = dedupeEnv(environment) // One provenance list covers every Kubernetes input, including hashed Secret @@ -445,6 +439,17 @@ func agentTemplateCard(template *v1alpha3.AgentTemplate) *a2atype.AgentCard { // dedupeEnv preserves first-seen ordering but gives the last value for a name // precedence, matching how compiler layers are applied. +func applyOperatorOnlyEnv(values []corev1.EnvVar, variable env.StringVar) []corev1.EnvVar { + name := variable.Name() + values = slices.DeleteFunc(values, func(item corev1.EnvVar) bool { + return item.Name == name + }) + if value := variable.Get(); value != "" { + values = append(values, corev1.EnvVar{Name: name, Value: value}) + } + return values +} + func dedupeEnv(values []corev1.EnvVar) []corev1.EnvVar { result := make([]corev1.EnvVar, 0, len(values)) index := map[string]int{} diff --git a/helm/kagent/templates/_helpers.tpl b/helm/kagent/templates/_helpers.tpl index 0448597778..02c0468e27 100644 --- a/helm/kagent/templates/_helpers.tpl +++ b/helm/kagent/templates/_helpers.tpl @@ -284,3 +284,20 @@ imagePullSecrets: {{- toYaml $global | nindent 2 }} {{- end -}} {{- end -}} + +{{/* +Serialize otel.tracing.contextKeys for KAGENT_TRACE_CONTEXT_KEYS. +A list of strings is joined with commas; any mapping entry is emitted as JSON +so {from, to, hash} objects survive into the runtime allowlist parser. +*/}} +{{- define "kagent.traceContextKeys" -}} +{{- $needsJSON := false -}} +{{- range . -}} +{{- if kindIs "map" . -}}{{- $needsJSON = true -}}{{- end -}} +{{- end -}} +{{- if $needsJSON -}} +{{- . | toJson -}} +{{- else -}} +{{- join "," . -}} +{{- end -}} +{{- end -}} diff --git a/helm/kagent/templates/controller-configmap.yaml b/helm/kagent/templates/controller-configmap.yaml index 32a66c833c..c4a93c89d0 100644 --- a/helm/kagent/templates/controller-configmap.yaml +++ b/helm/kagent/templates/controller-configmap.yaml @@ -21,7 +21,7 @@ data: OTEL_TRACING_ENABLED: {{ .Values.otel.tracing.enabled | quote }} OTEL_LOGGING_ENABLED: {{ .Values.otel.logging.enabled | quote }} {{- with .Values.otel.tracing.contextKeys }} - KAGENT_TRACE_CONTEXT_KEYS: {{ join "," . | quote }} + KAGENT_TRACE_CONTEXT_KEYS: {{ include "kagent.traceContextKeys" . | quote }} {{- end }} {{- $tracesEndpoint := .Values.otel.tracing.exporter.otlp.endpoint }} {{- $logsEndpoint := .Values.otel.logging.exporter.otlp.endpoint }} diff --git a/helm/kagent/templates/controller-deployment.yaml b/helm/kagent/templates/controller-deployment.yaml index 7e74c94378..d111719dd3 100644 --- a/helm/kagent/templates/controller-deployment.yaml +++ b/helm/kagent/templates/controller-deployment.yaml @@ -141,6 +141,15 @@ spec: {{- with .Values.controller.env }} {{- toYaml . | nindent 12 }} {{- end }} + {{- with .Values.otel.tracing.contextHashKeySecret }} + {{- if .name }} + - name: KAGENT_TRACE_CONTEXT_HASH_KEY + valueFrom: + secretKeyRef: + name: {{ .name | quote }} + key: {{ .key | default "hmac-key" | quote }} + {{- end }} + {{- end }} {{- if and .Values.controller.substrate .Values.controller.substrate.enabled }} - name: SUBSTRATE_ATE_API_ENDPOINT value: {{ .Values.controller.substrate.ateApiEndpoint | quote }} diff --git a/helm/kagent/tests/controller-deployment_test.yaml b/helm/kagent/tests/controller-deployment_test.yaml index c94fdf3785..b2f3ad4b54 100644 --- a/helm/kagent/tests/controller-deployment_test.yaml +++ b/helm/kagent/tests/controller-deployment_test.yaml @@ -842,10 +842,48 @@ tests: template: controller-configmap.yaml set: otel.tracing.contextKeys: - - user.email - - user.name + - sub - thread_id + - channel asserts: - equal: path: data.KAGENT_TRACE_CONTEXT_KEYS - value: "user.email,user.name,thread_id" + value: "sub,thread_id,channel" + + - it: should encode contextKeys mappings as JSON + template: controller-configmap.yaml + set: + otel.tracing.contextKeys: + - from: sub + to: user.id + - from: thread_id + to: kagent.thread_id + - channel + asserts: + - equal: + path: data.KAGENT_TRACE_CONTEXT_KEYS + value: '[{"from":"sub","to":"user.id"},{"from":"thread_id","to":"kagent.thread_id"},"channel"]' + + - it: should omit KAGENT_TRACE_CONTEXT_HASH_KEY by default + template: controller-deployment.yaml + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: KAGENT_TRACE_CONTEXT_HASH_KEY + + - it: should inject KAGENT_TRACE_CONTEXT_HASH_KEY from the configured secret + template: controller-deployment.yaml + set: + otel.tracing.contextHashKeySecret: + name: trace-context-hmac + key: hmac-key + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: KAGENT_TRACE_CONTEXT_HASH_KEY + valueFrom: + secretKeyRef: + name: trace-context-hmac + key: hmac-key diff --git a/helm/kagent/values.yaml b/helm/kagent/values.yaml index 49074317c4..31942c1300 100644 --- a/helm/kagent/values.yaml +++ b/helm/kagent/values.yaml @@ -848,11 +848,31 @@ oauth2-proxy: otel: tracing: enabled: false - # Allowlist of caller-supplied context keys promoted onto every agent span - # as kagent.context.. Values are read from W3C baggage and A2A message - # metadata. Empty (the default) disables promotion. - # e.g. ["user.email", "user.name", "thread_id", "channel"] + # Allowlist of caller-supplied context keys promoted onto every agent span. + # Values are read from W3C baggage and A2A message metadata. Empty (the + # default) disables promotion. + # + # Entries may be a source key or a mapping. Prefer an opaque identifier + # such as an OIDC subject for user.id; do not put names or email addresses + # on spans. Registry names (user.*, enduser.*, session.id) are left + # unprefixed; names already in the kagent. namespace are left as-is; + # everything else is emitted as kagent.context.. + # + # contextKeys: + # - {from: sub, to: user.id} + # - {from: thread_id, to: kagent.thread_id} + # - channel + # + # To derive a stable identifier without putting the original value on the + # span, set hash: hmac-sha256 and provide contextHashKeySecret: + # - {from: email, to: user.hash, hash: hmac-sha256} contextKeys: [] + # Secret providing KAGENT_TRACE_CONTEXT_HASH_KEY for hash: hmac-sha256 + # mappings. Injected into the controller and forwarded to agent runtimes. + # A Harness cannot set or replace it. + contextHashKeySecret: + name: "" + key: hmac-key exporter: otlp: endpoint: "" diff --git a/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py b/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py index 2956217367..67853b3945 100644 --- a/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py +++ b/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py @@ -1,25 +1,49 @@ """Promote allowlisted caller context onto every span of an agent request.""" +from __future__ import annotations + +import hashlib +import hmac +import json import os +from dataclasses import dataclass from typing import Any, Optional from opentelemetry import baggage from opentelemetry import context as otel_context -# Comma-separated allowlist of caller-supplied context keys to promote onto -# agent spans. Unset or empty (the default) disables promotion entirely. +# Allowlist of caller-supplied context keys to promote onto agent spans. +# Unset or empty (the default) disables promotion entirely. +# +# Accepts a comma-separated list of source keys, or a JSON array of strings +# and {from, to, hash} objects. TRACE_CONTEXT_KEYS_ENV_VAR = "KAGENT_TRACE_CONTEXT_KEYS" -# Namespaces every promoted value. Because the prefix is applied -# unconditionally, caller-supplied data cannot shadow a semantic convention -# attribute such as ``service.name``. +# HMAC key used when a mapping sets hash: hmac-sha256. Required for those +# entries; without it the hashed attribute is skipped rather than emitted +# in plaintext. +TRACE_CONTEXT_HASH_KEY_ENV_VAR = "KAGENT_TRACE_CONTEXT_HASH_KEY" + +# Namespaces custom promoted values so they cannot shadow a semantic +# convention attribute such as ``service.name``. Registry names +# (``user.*``, ``enduser.*``, ``session.id``) and names already in the +# ``kagent.`` namespace are left unprefixed; see ``_span_attribute_name``. CONTEXT_ATTRIBUTE_PREFIX = "kagent.context." +HASH_HMAC_SHA256 = "hmac-sha256" + MAX_CONTEXT_KEYS = 32 MAX_CONTEXT_KEY_LENGTH = 64 MAX_CONTEXT_VALUE_LENGTH = 256 +@dataclass(frozen=True) +class _ContextMapping: + source: str + attribute: str + hash: str = "" + + def caller_context_attributes( metadata: Optional[dict[str, Any]] = None, context: Optional[otel_context.Context] = None, @@ -35,35 +59,44 @@ def caller_context_attributes( metadata, which is the more specific source for a single message and therefore wins. Both are untrusted input, so keys must appear in the allowlist, values are stripped of control characters and truncated, and - every attribute is namespaced under ``CONTEXT_ATTRIBUTE_PREFIX``. + attribute names go through ``_span_attribute_name``. Args: metadata: A2A message metadata as a plain dict (may be ``None``). context: OTel context to read baggage from. Defaults to the current one. Returns: - Prefixed attribute name to sanitised value. Empty when the allowlist is - empty, which is the default. + Attribute name to sanitised value. Empty when the allowlist is empty, + which is the default. """ - keys = _allowed_context_keys() - if not keys: + mappings = _allowed_context_mappings() + if not mappings: return {} bag = baggage.get_all(context) attributes: dict[str, str] = {} - for key in keys: - value = _sanitize_context_value(bag.get(key)) + for mapping in mappings: + value = _sanitize_context_value(bag.get(mapping.source)) if metadata is not None: - scalar = _scalar_string(metadata.get(key)) + scalar = _scalar_string(metadata.get(mapping.source)) if scalar is not None: value = _sanitize_context_value(scalar) if not value: continue - attributes[CONTEXT_ATTRIBUTE_PREFIX + key] = value + if mapping.hash: + value = _hash_context_value(value, mapping.hash) + if not value: + continue + else: + value = _truncate_context_value(value) + name = _span_attribute_name(mapping.attribute) + if name in attributes: + continue + attributes[name] = value return attributes -def _allowed_context_keys() -> list[str]: +def _allowed_context_mappings() -> list[_ContextMapping]: """Parse the ``KAGENT_TRACE_CONTEXT_KEYS`` allowlist. Keys that are empty, over-long, or contain whitespace or control characters @@ -73,18 +106,93 @@ def _allowed_context_keys() -> list[str]: raw = os.getenv(TRACE_CONTEXT_KEYS_ENV_VAR, "").strip() if not raw: return [] + if raw.startswith("["): + return _cap_mappings(_parse_json_allowlist(raw)) + return _cap_mappings(_parse_comma_allowlist(raw)) + - keys: list[str] = [] +def _parse_comma_allowlist(raw: str) -> list[_ContextMapping]: + mappings: list[_ContextMapping] = [] for candidate in raw.split(","): - key = candidate.strip() - if not key or len(key) > MAX_CONTEXT_KEY_LENGTH or not _is_attribute_key(key): - continue - if key in keys: + mapping = _new_context_mapping(candidate.strip(), "", "") + if mapping is not None: + mappings.append(mapping) + return mappings + + +def _parse_json_allowlist(raw: str) -> list[_ContextMapping]: + try: + items = json.loads(raw) + except json.JSONDecodeError: + return [] + if not isinstance(items, list): + return [] + mappings: list[_ContextMapping] = [] + for item in items: + if isinstance(item, str): + mapping = _new_context_mapping(item, "", "") + elif isinstance(item, dict): + from_key, to_key, hash_alg = item.get("from"), item.get("to"), item.get("hash") + if from_key is not None and not isinstance(from_key, str): + mapping = None + else: + mapping = _new_context_mapping( + from_key or "", + to_key if isinstance(to_key, str) else "", + hash_alg if isinstance(hash_alg, str) else "", + ) + else: + mapping = None + if mapping is not None: + mappings.append(mapping) + return mappings + + +def _new_context_mapping(from_key: str, to_key: str, hash_alg: str) -> Optional[_ContextMapping]: + from_key = from_key.strip() + to_key = to_key.strip() + hash_alg = hash_alg.strip() + if not from_key or len(from_key) > MAX_CONTEXT_KEY_LENGTH or not _is_attribute_key(from_key): + return None + if not to_key: + to_key = from_key + if len(to_key) > MAX_CONTEXT_KEY_LENGTH or not _is_attribute_key(to_key): + return None + if hash_alg and hash_alg != HASH_HMAC_SHA256: + return None + return _ContextMapping(source=from_key, attribute=to_key, hash=hash_alg) + + +def _cap_mappings(mappings: list[_ContextMapping]) -> list[_ContextMapping]: + out: list[_ContextMapping] = [] + seen: set[tuple[str, str, str]] = set() + for mapping in mappings: + identity = (mapping.source, mapping.attribute, mapping.hash) + if identity in seen: continue - keys.append(key) - if len(keys) == MAX_CONTEXT_KEYS: + seen.add(identity) + out.append(mapping) + if len(out) == MAX_CONTEXT_KEYS: break - return keys + return out + + +def _span_attribute_name(name: str) -> str: + """Return the name written onto the span. + + ``user.*``, ``enduser.*``, and ``session.id`` pass through unprefixed so + operators can use the semantic convention names. Names already in the + ``kagent.`` namespace are left as-is. Everything else is placed under + ``kagent.context.`` so a caller-supplied ``service.name`` cannot shadow + the real one. + """ + if _is_registry_attribute(name) or name.startswith("kagent."): + return name + return CONTEXT_ATTRIBUTE_PREFIX + name + + +def _is_registry_attribute(name: str) -> bool: + return name.startswith("user.") or name.startswith("enduser.") or name == "session.id" def _is_attribute_key(key: str) -> bool: @@ -119,13 +227,25 @@ def _scalar_string(value: Any) -> Optional[str]: def _sanitize_context_value(value: Any) -> str: - """Make an untrusted value safe to attach to a span. + """Drop control characters and trim space so a value cannot forge structure.""" + if not isinstance(value, str): + return "" + return "".join(char for char in value if not _is_control(char)).strip() - Control characters are dropped so a value cannot forge structure in a - downstream trace or log renderer, and the result is truncated to bound the - size of exported spans. + +def _truncate_context_value(value: str) -> str: + return value[:MAX_CONTEXT_VALUE_LENGTH] + + +def _hash_context_value(value: str, algorithm: str) -> str: + """Hash *value* with the requested algorithm. + + Unknown algorithms and a missing HMAC key skip the attribute: never fall + back to putting the original value on the span. """ - if not isinstance(value, str): + if algorithm != HASH_HMAC_SHA256: + return "" + key = os.getenv(TRACE_CONTEXT_HASH_KEY_ENV_VAR, "") + if not key: return "" - cleaned = "".join(char for char in value if not _is_control(char)).strip() - return cleaned[:MAX_CONTEXT_VALUE_LENGTH] + return hmac.new(key.encode("utf-8"), value.encode("utf-8"), hashlib.sha256).hexdigest() diff --git a/python/packages/kagent-core/tests/test_caller_context_attributes.py b/python/packages/kagent-core/tests/test_caller_context_attributes.py index c82b9311e9..2584f84ccf 100644 --- a/python/packages/kagent-core/tests/test_caller_context_attributes.py +++ b/python/packages/kagent-core/tests/test_caller_context_attributes.py @@ -1,3 +1,6 @@ +import hashlib +import hmac + import pytest from opentelemetry import baggage from opentelemetry import context as otel_context @@ -11,8 +14,9 @@ MAX_CONTEXT_KEY_LENGTH, MAX_CONTEXT_KEYS, MAX_CONTEXT_VALUE_LENGTH, + TRACE_CONTEXT_HASH_KEY_ENV_VAR, TRACE_CONTEXT_KEYS_ENV_VAR, - _allowed_context_keys, + _allowed_context_mappings, ) from kagent.core.tracing._span_processor import ( KagentAttributesSpanProcessor, @@ -28,26 +32,28 @@ def baggage_context(members: dict[str, str]) -> otel_context.Context: return context +def hmac_sha256_hex(key: str, value: str) -> str: + return hmac.new(key.encode("utf-8"), value.encode("utf-8"), hashlib.sha256).hexdigest() + + class TestCallerContextAttributes: """Tests for allowlist-driven promotion of caller context onto spans.""" - def test_disabled_by_default(self, monkeypatch): + def test_empty_allowlist_disables_promotion(self, monkeypatch): monkeypatch.delenv(TRACE_CONTEXT_KEYS_ENV_VAR, raising=False) assert ( caller_context_attributes( {"thread_id": "T1"}, - baggage_context({"user.email": "ada@example.com"}), + baggage_context({"sub": "opaque-subject"}), ) == {} ) def test_promotes_allowlisted_baggage(self, monkeypatch): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "user.email,user.name") - assert caller_context_attributes( - None, baggage_context({"user.email": "ada@example.com", "user.name": "Ada"}) - ) == { - "kagent.context.user.email": "ada@example.com", - "kagent.context.user.name": "Ada", + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "sub,thread_id") + assert caller_context_attributes(None, baggage_context({"sub": "opaque-subject", "thread_id": "T123"})) == { + "kagent.context.sub": "opaque-subject", + "kagent.context.thread_id": "T123", } def test_promotes_allowlisted_message_metadata(self, monkeypatch): @@ -58,16 +64,16 @@ def test_promotes_allowlisted_message_metadata(self, monkeypatch): } def test_message_metadata_overrides_baggage(self, monkeypatch): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "user.email") + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "sub") assert caller_context_attributes( - {"user.email": "from-metadata@example.com"}, - baggage_context({"user.email": "from-baggage@example.com"}), - ) == {"kagent.context.user.email": "from-metadata@example.com"} + {"sub": "from-metadata"}, + baggage_context({"sub": "from-baggage"}), + ) == {"kagent.context.sub": "from-metadata"} def test_ignores_keys_outside_the_allowlist(self, monkeypatch): monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "thread_id") assert caller_context_attributes( - {"thread_id": "T1", "customer.pan": "4111111111111111"}, + {"thread_id": "T1", "extra": "nope"}, baggage_context({"secret.token": "s3cret"}), ) == {"kagent.context.thread_id": "T1"} @@ -103,32 +109,99 @@ def test_truncates_long_values(self, monkeypatch): assert len(promoted["kagent.context.note"]) == MAX_CONTEXT_VALUE_LENGTH def test_cannot_shadow_semantic_conventions(self, monkeypatch): - """The prefix is what stops caller data replacing service.name and friends.""" + """Custom keys still cannot replace service.name. Registry names are the exception.""" monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "service.name") promoted = caller_context_attributes({"service.name": "impostor"}) assert "service.name" not in promoted assert promoted == {f"{CONTEXT_ATTRIBUTE_PREFIX}service.name": "impostor"} + def test_registry_attributes_stay_unprefixed(self, monkeypatch): + """user.*, enduser.*, and session.id are semantic convention names.""" + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "user.id,enduser.id,session.id,channel") + assert caller_context_attributes( + { + "user.id": "opaque-subject", + "enduser.id": "end-user", + "session.id": "sess-1", + "channel": "C0AB1", + } + ) == { + "user.id": "opaque-subject", + "enduser.id": "end-user", + "session.id": "sess-1", + "kagent.context.channel": "C0AB1", + } + + def test_session_id_is_unprefixed_but_session_foo_is_not(self, monkeypatch): + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "session.id,session.foo") + assert caller_context_attributes({"session.id": "sess-1", "session.foo": "other"}) == { + "session.id": "sess-1", + "kagent.context.session.foo": "other", + } + + def test_maps_source_keys_onto_registry_and_kagent_names(self, monkeypatch): + monkeypatch.setenv( + TRACE_CONTEXT_KEYS_ENV_VAR, + '[{"from":"sub","to":"user.id"},{"from":"thread_id","to":"kagent.thread_id"},"channel"]', + ) + assert caller_context_attributes({"sub": "opaque-subject", "thread_id": "T123", "channel": "C0AB1"}) == { + "user.id": "opaque-subject", + "kagent.thread_id": "T123", + "kagent.context.channel": "C0AB1", + } + + def test_invalid_json_allowlist_promotes_nothing(self, monkeypatch): + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, '[{"from":"sub"') + assert caller_context_attributes({"sub": "opaque-subject"}) == {} + + def test_hashes_with_hmac_sha256(self, monkeypatch): + key = "test-hmac-key" + monkeypatch.setenv( + TRACE_CONTEXT_KEYS_ENV_VAR, + '[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]', + ) + monkeypatch.setenv(TRACE_CONTEXT_HASH_KEY_ENV_VAR, key) + promoted = caller_context_attributes({"email": "ada@example.com"}) + assert promoted == {"user.hash": hmac_sha256_hex(key, "ada@example.com")} + assert not any("@example.com" in value for value in promoted.values()) + + def test_hash_without_key_emits_nothing(self, monkeypatch): + """Missing HMAC key must not fall back to putting the original value on the span.""" + monkeypatch.setenv( + TRACE_CONTEXT_KEYS_ENV_VAR, + '[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]', + ) + monkeypatch.delenv(TRACE_CONTEXT_HASH_KEY_ENV_VAR, raising=False) + assert caller_context_attributes({"email": "ada@example.com"}) == {} + + def test_unknown_hash_emits_nothing(self, monkeypatch): + monkeypatch.setenv( + TRACE_CONTEXT_KEYS_ENV_VAR, + '[{"from":"email","to":"user.hash","hash":"md5"}]', + ) + monkeypatch.setenv(TRACE_CONTEXT_HASH_KEY_ENV_VAR, "test-hmac-key") + assert caller_context_attributes({"email": "ada@example.com"}) == {} + -class TestAllowedContextKeys: +class TestAllowedContextMappings: """Tests for allowlist parsing and its bounds.""" def test_caps_list_length(self, monkeypatch): keys = ",".join(f"key{index}" for index in range(MAX_CONTEXT_KEYS * 2)) monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, keys) - assert len(_allowed_context_keys()) == MAX_CONTEXT_KEYS + assert len(_allowed_context_mappings()) == MAX_CONTEXT_KEYS def test_drops_over_long_and_duplicate_keys(self, monkeypatch): monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, f"a,a,{'b' * (MAX_CONTEXT_KEY_LENGTH + 1)},c") - assert _allowed_context_keys() == ["a", "c"] + assert [mapping.source for mapping in _allowed_context_mappings()] == ["a", "c"] def test_drops_keys_that_are_not_valid_attribute_names(self, monkeypatch): monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "good, bad key ,\tanother\tbad") - assert _allowed_context_keys() == ["good"] + assert [mapping.source for mapping in _allowed_context_mappings()] == ["good"] def test_empty_allowlist_disables_promotion(self, monkeypatch): monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, " , ,") - assert _allowed_context_keys() == [] + assert _allowed_context_mappings() == [] class TestPromotedAttributesReachEverySpan: @@ -156,20 +229,21 @@ def record_spans(span_attributes: dict) -> dict[str, dict]: return {span.name: dict(span.attributes or {}) for span in exporter.get_finished_spans()} def test_flag_on_stamps_every_span(self, monkeypatch): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "user.email") - promoted = caller_context_attributes(None, baggage_context({"user.email": "ada@example.com"})) + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, '[{"from":"sub","to":"user.id"}]') + promoted = caller_context_attributes(None, baggage_context({"sub": "opaque-subject"})) spans = self.record_spans(promoted) assert set(spans) == {"root", "execute_tool", "generate_content"} for attributes in spans.values(): - assert attributes["kagent.context.user.email"] == "ada@example.com" + assert attributes["user.id"] == "opaque-subject" def test_flag_off_leaves_spans_unchanged(self, monkeypatch): monkeypatch.delenv(TRACE_CONTEXT_KEYS_ENV_VAR, raising=False) - promoted = caller_context_attributes({"thread_id": "T1"}, baggage_context({"user.email": "ada@example.com"})) + promoted = caller_context_attributes({"thread_id": "T1"}, baggage_context({"sub": "opaque-subject"})) spans = self.record_spans(promoted) for attributes in spans.values(): + assert "user.id" not in attributes assert not [key for key in attributes if key.startswith(CONTEXT_ATTRIBUTE_PREFIX)] diff --git a/python/packages/kagent-core/tests/test_tracing_configure.py b/python/packages/kagent-core/tests/test_tracing_configure.py index da9f0ac810..b1415c66d8 100644 --- a/python/packages/kagent-core/tests/test_tracing_configure.py +++ b/python/packages/kagent-core/tests/test_tracing_configure.py @@ -208,9 +208,9 @@ def test_otel_sdk_default_propagator_includes_baggage(): baggage from the default propagator, this test will fail and explicit configuration will be needed. """ - ctx = get_global_textmap().extract({"baggage": "user.email=ada%40example.com"}) + ctx = get_global_textmap().extract({"baggage": "sub=opaque-subject"}) - assert get_baggage("user.email", ctx) == "ada@example.com" + assert get_baggage("sub", ctx) == "opaque-subject" @pytest.mark.parametrize( From 9de37b831a554fa13b349420d3d3c30c33fc2ff4 Mon Sep 17 00:00:00 2001 From: Ricardo Temperini <29879569+rtemperini@users.noreply.github.com> Date: Mon, 31 Aug 2026 10:51:25 +0200 Subject: [PATCH 06/13] fix: fill caller context only when unset and tighten names Caller-supplied attributes no longer override runtime keys. Registry passthrough is an explicit set, empty metadata cannot wipe baggage, and the allowlist is parsed once per process. Signed-off-by: Ricardo Temperini <29879569+rtemperini@users.noreply.github.com> --- docs/architecture/trace-context.md | 28 ++-- go/adk/pkg/a2a/executor.go | 5 +- go/adk/pkg/telemetry/context_attributes.go | 109 +++++++++------- .../pkg/telemetry/context_attributes_test.go | 120 +++++++++++++++--- go/core/pkg/env/otel.go | 6 +- .../tests/controller-deployment_test.yaml | 4 +- helm/kagent/values.yaml | 7 +- .../src/kagent/adk/_agent_executor.py | 7 +- .../src/kagent/core/tracing/__init__.py | 4 +- .../core/tracing/_context_attributes.py | 88 +++++++++---- python/packages/kagent-core/tests/conftest.py | 11 ++ .../tests/test_caller_context_attributes.py | 96 +++++++++++++- .../src/kagent/crewai/_executor.py | 7 +- .../src/kagent/langgraph/_executor.py | 6 +- 14 files changed, 373 insertions(+), 125 deletions(-) create mode 100644 python/packages/kagent-core/tests/conftest.py diff --git a/docs/architecture/trace-context.md b/docs/architecture/trace-context.md index aed95a69a8..10c45c8b54 100644 --- a/docs/architecture/trace-context.md +++ b/docs/architecture/trace-context.md @@ -25,7 +25,7 @@ otel: enabled: true contextKeys: - {from: sub, to: user.id} - - {from: thread_id, to: kagent.thread_id} + - thread_id - channel ``` @@ -60,7 +60,9 @@ further plumbing. **A2A `message.metadata` is the complement** for callers that cannot set a header — for example, a bot that speaks A2A over an SDK that exposes message metadata but not transport headers. It is scoped to a single message, and -because it is the more specific source, it overrides baggage for the same key. +because it is the more specific source, it overrides baggage for the same key +when the sanitised metadata value is non-empty. A blank metadata entry cannot +wipe a baggage value. A key absent from both sources is simply not emitted. @@ -72,17 +74,21 @@ Each mapping is read from `from` (defaulting to the entry itself) and written as span attribute `to` (defaulting to `from`) after the prefix rules below: ``` -baggage: sub=opaque-subject → user.id = "opaque-subject" -metadata: {"thread_id": "1717171.42"} → kagent.thread_id = "1717171.42" -metadata: {"channel": "C0AB1"} → kagent.context.channel = "C0AB1" +baggage: sub=opaque-subject → user.id = "opaque-subject" +metadata: {"thread_id": "1717171.42"} → kagent.context.thread_id = "1717171.42" +metadata: {"channel": "C0AB1"} → kagent.context.channel = "C0AB1" ``` | Destination name | Emitted as | |---|---| -| `user.*`, `enduser.*`, `session.id` | Unprefixed (OpenTelemetry semantic conventions) | -| Already in the `kagent.` namespace | Unprefixed | +| `user.id`, `user.hash`, `enduser.id`, `session.id` | Unprefixed (OpenTelemetry semantic conventions) | | Anything else | `kagent.context.` | +Invented names such as `user.asdasd` or `kagent.user_id` are prefixed. Promoted +values fill the request-scoped bag only when the key is not already set, so +caller context cannot override `kagent.user_id`, `kagent.app_name`, or +`gen_ai.*` attributes the runtime stamped. + The attributes are merged into the **request-scoped attribute bag**, not set on a single span. The `KagentAttributesSpanProcessor` (Python) and `kagentAttributesSpanProcessor` (Go) stamp that bag onto every span started @@ -109,7 +115,7 @@ appear on a span, hash it with HMAC-SHA256 onto the registry attribute contextKeys: - {from: sub, to: user.id} - {from: email, to: user.hash, hash: hmac-sha256} - - {from: thread_id, to: kagent.thread_id} + - thread_id ``` `hash: hmac-sha256` requires `KAGENT_TRACE_CONTEXT_HASH_KEY` (Helm: @@ -134,7 +140,7 @@ axis: | Attribute explosion / cardinality | Only allowlisted keys are read; the allowlist itself is capped at 32 entries | | Oversized spans | Values are truncated to 256 characters, keys to 64 | | Log or trace injection | Control characters are stripped from values | -| Shadowing semantic conventions | Custom keys are namespaced under `kagent.context.`; only `user.*`, `enduser.*`, and `session.id` pass through unprefixed | +| Shadowing semantic conventions | Custom keys are namespaced under `kagent.context.`; only `user.id`, `user.hash`, `enduser.id`, and `session.id` pass through unprefixed. Caller context fills missing keys only. | | Leaking secrets into a trace backend | Nothing is promoted unless an operator names the key; hashed entries are omitted when the HMAC key is unset | | A tenant widening the allowlist | The allowlist is cluster-wide operator configuration; an entry of the same name in a `Harness` environment is dropped rather than inherited | @@ -156,8 +162,8 @@ already sits between kagent and the backend: processors: transform: trace_statements: - - set(span.attributes["session.id"], span.attributes["kagent.thread_id"]) - where span.attributes["kagent.thread_id"] != nil + - set(span.attributes["session.id"], span.attributes["kagent.context.thread_id"]) + where span.attributes["kagent.context.thread_id"] != nil ``` --- diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index 2201806485..10c48628ff 100644 --- a/go/adk/pkg/a2a/executor.go +++ b/go/adk/pkg/a2a/executor.go @@ -4,7 +4,6 @@ import ( "context" "fmt" "iter" - "maps" "strings" a2atype "github.com/a2aproject/a2a-go/v2/a2a" @@ -132,7 +131,9 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.ExecutorCon } // Allowlisted caller context joins the request-scoped bag rather than a // single span, so tool, sub-agent, and model spans all carry it. - maps.Copy(spanAttributes, telemetry.CallerContextAttributes(ctx, reqCtx.Message.Metadata)) + // Fill-if-absent so a caller cannot override kagent.user_id or + // gen_ai.* attributes the runtime already stamped. + telemetry.MergeCallerContextAttributes(spanAttributes, ctx, reqCtx.Message.Metadata) ctx = telemetry.SetKAgentSpanAttributes(ctx, spanAttributes) ctx, invocationSpan := telemetry.StartInvocationSpan(ctx) defer invocationSpan.End() diff --git a/go/adk/pkg/telemetry/context_attributes.go b/go/adk/pkg/telemetry/context_attributes.go index d0774d144f..145c3e7689 100644 --- a/go/adk/pkg/telemetry/context_attributes.go +++ b/go/adk/pkg/telemetry/context_attributes.go @@ -9,29 +9,18 @@ import ( "os" "strconv" "strings" + "sync" "unicode" + "github.com/kagent-dev/kagent/go/core/pkg/env" "go.opentelemetry.io/otel/baggage" ) const ( - // traceContextKeysEnvVar holds the allowlist of caller-supplied context - // keys to promote onto agent spans. Unset or empty (the default) disables - // promotion entirely. - // - // Accepts a comma-separated list of source keys, or a JSON array of strings - // and {from, to, hash} objects. See allowedContextMappings. - traceContextKeysEnvVar = "KAGENT_TRACE_CONTEXT_KEYS" - - // traceContextHashKeyEnvVar is the HMAC key used when a mapping sets - // hash: hmac-sha256. Required for those entries; without it the hashed - // attribute is skipped rather than emitted in plaintext. - traceContextHashKeyEnvVar = "KAGENT_TRACE_CONTEXT_HASH_KEY" - // contextAttributePrefix namespaces custom promoted values so they cannot - // shadow a semantic convention attribute such as service.name. Registry - // names (user.*, enduser.*, session.id) and names already in the kagent. - // namespace are left unprefixed; see spanAttributeName. + // shadow a semantic convention attribute such as service.name. Only the + // explicit registry names in unprefixedAttributes pass through as-is; + // see spanAttributeName. contextAttributePrefix = "kagent.context." hashHMACSHA256 = "hmac-sha256" @@ -41,6 +30,23 @@ const ( maxContextValueLength = 256 ) +// Names come from env.KagentTraceContextKeys / HashKey so the controller and +// the runtime share one source of truth. +var ( + traceContextKeysEnvVar = env.KagentTraceContextKeys.Name() + traceContextHashKeyEnvVar = env.KagentTraceContextHashKey.Name() +) + +// unprefixedAttributes is the OpenTelemetry registry names callers may emit +// without a kagent.context. prefix. It is an explicit set so an invented +// name such as user.asdasd cannot masquerade as a semantic convention. +var unprefixedAttributes = map[string]struct{}{ + "user.id": {}, + "user.hash": {}, + "enduser.id": {}, + "session.id": {}, +} + // contextMapping is one allowlisted promotion: read source from baggage or // A2A metadata and emit it as span attribute attribute (after prefix rules). type contextMapping struct { @@ -49,6 +55,19 @@ type contextMapping struct { hash string } +// loadAllowedContextMappings parses KAGENT_TRACE_CONTEXT_KEYS once. The +// allowlist cannot change without a process restart; tests call +// resetAllowedContextMappings after t.Setenv. +var loadAllowedContextMappings = sync.OnceValue(parseAllowedContextMappings) + +func resetAllowedContextMappings() { + loadAllowedContextMappings = sync.OnceValue(parseAllowedContextMappings) +} + +func allowedContextMappings() []contextMapping { + return loadAllowedContextMappings() +} + // CallerContextAttributes returns the caller-supplied context values that an // operator allowlisted through KAGENT_TRACE_CONTEXT_KEYS. Merge the result into // the request-scoped attribute bag (see SetKAgentSpanAttributes) so every span @@ -57,9 +76,10 @@ type contextMapping struct { // // Values are read from W3C baggage first and then from the A2A message // metadata, which is the more specific source for a single message and -// therefore wins. Both are untrusted input, so keys must appear in the -// allowlist, values are stripped of control characters and truncated, and -// attribute names go through spanAttributeName. +// therefore wins — but only when the sanitised metadata scalar is non-empty, +// so a blank metadata entry cannot wipe a baggage value. Both are untrusted +// input, so keys must appear in the allowlist, values are stripped of control +// characters and truncated, and attribute names go through spanAttributeName. // // Returns nil when the allowlist is empty, which is the default. func CallerContextAttributes(ctx context.Context, metadata map[string]any) map[string]string { @@ -73,7 +93,9 @@ func CallerContextAttributes(ctx context.Context, metadata map[string]any) map[s for _, mapping := range mappings { value := sanitizeContextValue(bag.Member(mapping.source).Value()) if scalar, ok := scalarString(metadata[mapping.source]); ok { - value = sanitizeContextValue(scalar) + if cleaned := sanitizeContextValue(scalar); cleaned != "" { + value = cleaned + } } if value == "" { continue @@ -98,11 +120,19 @@ func CallerContextAttributes(ctx context.Context, metadata map[string]any) map[s return attrs } -// allowedContextMappings parses the KAGENT_TRACE_CONTEXT_KEYS allowlist. -// Entries that are empty, over-long, or contain whitespace or control -// characters are dropped, and the list is capped at maxContextKeys so a -// misconfigured allowlist cannot inflate span cardinality without bound. -func allowedContextMappings() []contextMapping { +// MergeCallerContextAttributes copies promoted caller context into dst without +// replacing keys that are already set, so caller-supplied data cannot override +// kagent.user_id, kagent.app_name, or other attributes the runtime already +// stamped. +func MergeCallerContextAttributes(dst map[string]string, ctx context.Context, metadata map[string]any) { + for key, value := range CallerContextAttributes(ctx, metadata) { + if _, exists := dst[key]; !exists { + dst[key] = value + } + } +} + +func parseAllowedContextMappings() []contextMapping { raw := strings.TrimSpace(os.Getenv(traceContextKeysEnvVar)) if raw == "" { return nil @@ -156,13 +186,13 @@ func newContextMapping(from, to, hash string) (contextMapping, bool) { from = strings.TrimSpace(from) to = strings.TrimSpace(to) hash = strings.TrimSpace(hash) - if from == "" || len(from) > maxContextKeyLength || !isAttributeKey(from) { + if from == "" || len([]rune(from)) > maxContextKeyLength || !isAttributeKey(from) { return contextMapping{}, false } if to == "" { to = from } - if len(to) > maxContextKeyLength || !isAttributeKey(to) { + if len([]rune(to)) > maxContextKeyLength || !isAttributeKey(to) { return contextMapping{}, false } if hash != "" && hash != hashHMACSHA256 { @@ -190,23 +220,17 @@ func capMappings(mappings []contextMapping) []contextMapping { // spanAttributeName is the name written onto the span. // -// user.*, enduser.*, and session.id pass through unprefixed so operators can -// use the semantic convention names. Names already in the kagent. namespace -// are left as-is. Everything else is placed under kagent.context. so a -// caller-supplied service.name cannot shadow the real one. +// user.id, user.hash, enduser.id, and session.id pass through unprefixed so +// operators can use the semantic convention names. Everything else is placed +// under kagent.context. so a caller-supplied service.name or kagent.user_id +// cannot shadow the real one. func spanAttributeName(name string) string { - if isRegistryAttribute(name) || strings.HasPrefix(name, "kagent.") { + if _, ok := unprefixedAttributes[name]; ok { return name } return contextAttributePrefix + name } -func isRegistryAttribute(name string) bool { - return strings.HasPrefix(name, "user.") || - strings.HasPrefix(name, "enduser.") || - name == "session.id" -} - // isAttributeKey reports whether key is safe to use as a span attribute name. func isAttributeKey(key string) bool { return strings.IndexFunc(key, func(r rune) bool { @@ -216,7 +240,8 @@ func isAttributeKey(key string) bool { // scalarString renders a JSON scalar from A2A message metadata as a string. // Objects and arrays are skipped: they are unbounded in size and carry no -// useful meaning as a span attribute value. +// useful meaning as a span attribute value. JSON and protobuf Struct decode +// numbers as float64, so that is the only numeric case. func scalarString(value any) (string, bool) { switch v := value.(type) { case string: @@ -224,11 +249,7 @@ func scalarString(value any) (string, bool) { case bool: return strconv.FormatBool(v), true case float64: - return strconv.FormatFloat(v, 'g', -1, 64), true - case int: - return strconv.Itoa(v), true - case int64: - return strconv.FormatInt(v, 10), true + return strconv.FormatFloat(v, 'f', -1, 64), true default: return "", false } diff --git a/go/adk/pkg/telemetry/context_attributes_test.go b/go/adk/pkg/telemetry/context_attributes_test.go index 89d9962184..c3cff09b7c 100644 --- a/go/adk/pkg/telemetry/context_attributes_test.go +++ b/go/adk/pkg/telemetry/context_attributes_test.go @@ -14,6 +14,13 @@ import ( "go.opentelemetry.io/otel/sdk/trace/tracetest" ) +func setAllowlist(t *testing.T, allowlist string) { + t.Helper() + t.Setenv(traceContextKeysEnvVar, allowlist) + resetAllowedContextMappings() + t.Cleanup(resetAllowedContextMappings) +} + func baggageContext(t *testing.T, members map[string]string) context.Context { t.Helper() @@ -80,6 +87,13 @@ func TestCallerContextAttributes(t *testing.T) { metadata: map[string]any{"sub": "from-metadata"}, want: map[string]string{"kagent.context.sub": "from-metadata"}, }, + { + name: "empty metadata does not override baggage", + allowlist: "sub", + baggageVals: map[string]string{"sub": "from-baggage"}, + metadata: map[string]any{"sub": " "}, + want: map[string]string{"kagent.context.sub": "from-baggage"}, + }, { name: "ignores keys outside the allowlist", allowlist: "thread_id", @@ -91,7 +105,7 @@ func TestCallerContextAttributes(t *testing.T) { name: "renders scalar metadata types", allowlist: "count,ratio,enabled", metadata: map[string]any{ - "count": int64(7), + "count": float64(7), "ratio": float64(2.5), "enabled": true, }, @@ -101,6 +115,12 @@ func TestCallerContextAttributes(t *testing.T) { "kagent.context.enabled": "true", }, }, + { + name: "formats large floats without scientific notation", + allowlist: "value", + metadata: map[string]any{"value": float64(1e20)}, + want: map[string]string{"kagent.context.value": "100000000000000000000"}, + }, { name: "skips non-scalar and empty metadata values", allowlist: "nested,list,blank", @@ -124,8 +144,7 @@ func TestCallerContextAttributes(t *testing.T) { want: map[string]string{"kagent.context.good": "yes"}, }, { - // Registry names pass through unprefixed so operators can use - // semantic convention attributes instead of inventing new ones. + // Only the explicit registry set passes through unprefixed. name: "registry attributes stay unprefixed", allowlist: "user.id,enduser.id,session.id,channel", metadata: map[string]any{ @@ -141,6 +160,15 @@ func TestCallerContextAttributes(t *testing.T) { "kagent.context.channel": "C0AB1", }, }, + { + name: "unknown user attributes are prefixed", + allowlist: `[{"from":"x","to":"user.asdasd"},{"from":"y","to":"user.email"}]`, + metadata: map[string]any{"x": "nope", "y": "also-nope"}, + want: map[string]string{ + "kagent.context.user.asdasd": "nope", + "kagent.context.user.email": "also-nope", + }, + }, { // session.id is the registry name; other session.* keys are not. name: "session.id is unprefixed but session.foo is not", @@ -152,8 +180,8 @@ func TestCallerContextAttributes(t *testing.T) { }, }, { - name: "maps source keys onto registry and kagent names", - allowlist: `[{"from":"sub","to":"user.id"},{"from":"thread_id","to":"kagent.thread_id"},"channel"]`, + name: "maps source keys onto registry names", + allowlist: `[{"from":"sub","to":"user.id"},{"from":"thread_id","to":"session.id"},"channel"]`, metadata: map[string]any{ "sub": "opaque-subject", "thread_id": "T123", @@ -161,10 +189,19 @@ func TestCallerContextAttributes(t *testing.T) { }, want: map[string]string{ "user.id": "opaque-subject", - "kagent.thread_id": "T123", + "session.id": "T123", "kagent.context.channel": "C0AB1", }, }, + { + name: "kagent destination names are prefixed", + allowlist: `[{"from":"uid","to":"kagent.user_id"},{"from":"tid","to":"kagent.thread_id"}]`, + metadata: map[string]any{"uid": "attacker", "tid": "T123"}, + want: map[string]string{ + "kagent.context.kagent.user_id": "attacker", + "kagent.context.kagent.thread_id": "T123", + }, + }, { name: "invalid JSON allowlist promotes nothing", allowlist: `[{"from":"sub"`, @@ -175,7 +212,7 @@ func TestCallerContextAttributes(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - t.Setenv(traceContextKeysEnvVar, tt.allowlist) + setAllowlist(t, tt.allowlist) if tt.hashKey != "" { t.Setenv(traceContextHashKeyEnvVar, tt.hashKey) } @@ -194,9 +231,32 @@ func TestCallerContextAttributes(t *testing.T) { } } +func TestMergeCallerContextAttributes_DoesNotOverrideExisting(t *testing.T) { + setAllowlist(t, `[{"from":"sub","to":"user.id"},"thread_id"]`) + + dst := map[string]string{ + "user.id": "runtime-user", + "kagent.user_id": "runtime-user", + } + MergeCallerContextAttributes(dst, context.Background(), map[string]any{ + "sub": "attacker", + "thread_id": "T1", + }) + + if dst["user.id"] != "runtime-user" { + t.Errorf("user.id = %q, want runtime-user", dst["user.id"]) + } + if dst["kagent.user_id"] != "runtime-user" { + t.Errorf("kagent.user_id = %q, want runtime-user", dst["kagent.user_id"]) + } + if dst["kagent.context.thread_id"] != "T1" { + t.Errorf("thread_id = %q, want T1", dst["kagent.context.thread_id"]) + } +} + func TestCallerContextAttributes_HashesWithHMACSHA256(t *testing.T) { const key = "test-hmac-key" - t.Setenv(traceContextKeysEnvVar, `[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]`) + setAllowlist(t, `[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]`) t.Setenv(traceContextHashKeyEnvVar, key) got := CallerContextAttributes(context.Background(), map[string]any{ @@ -216,7 +276,7 @@ func TestCallerContextAttributes_HashesWithHMACSHA256(t *testing.T) { func TestCallerContextAttributes_HashWithoutKeyEmitsNothing(t *testing.T) { // Missing HMAC key must not fall back to putting the original value on the span. - t.Setenv(traceContextKeysEnvVar, `[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]`) + setAllowlist(t, `[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]`) t.Setenv(traceContextHashKeyEnvVar, "") got := CallerContextAttributes(context.Background(), map[string]any{ @@ -228,7 +288,7 @@ func TestCallerContextAttributes_HashWithoutKeyEmitsNothing(t *testing.T) { } func TestCallerContextAttributes_UnknownHashEmitsNothing(t *testing.T) { - t.Setenv(traceContextKeysEnvVar, `[{"from":"email","to":"user.hash","hash":"md5"}]`) + setAllowlist(t, `[{"from":"email","to":"user.hash","hash":"md5"}]`) t.Setenv(traceContextHashKeyEnvVar, "test-hmac-key") got := CallerContextAttributes(context.Background(), map[string]any{ @@ -240,7 +300,7 @@ func TestCallerContextAttributes_UnknownHashEmitsNothing(t *testing.T) { } func TestCallerContextAttributes_TruncatesLongValues(t *testing.T) { - t.Setenv(traceContextKeysEnvVar, "note") + setAllowlist(t, "note") got := CallerContextAttributes(context.Background(), map[string]any{ "note": strings.Repeat("a", maxContextValueLength*2), @@ -256,7 +316,7 @@ func TestAllowedContextMappings_CapsListLength(t *testing.T) { for i := range maxContextKeys * 2 { keys = append(keys, "key"+strconv.Itoa(i)) } - t.Setenv(traceContextKeysEnvVar, strings.Join(keys, ",")) + setAllowlist(t, strings.Join(keys, ",")) if got := len(allowedContextMappings()); got != maxContextKeys { t.Errorf("allowlist length = %d, want %d", got, maxContextKeys) @@ -264,7 +324,7 @@ func TestAllowedContextMappings_CapsListLength(t *testing.T) { } func TestAllowedContextMappings_DropsOverLongAndDuplicateKeys(t *testing.T) { - t.Setenv(traceContextKeysEnvVar, "a,a,"+strings.Repeat("b", maxContextKeyLength+1)+",c") + setAllowlist(t, "a,a,"+strings.Repeat("b", maxContextKeyLength+1)+",c") got := allowedContextMappings() @@ -279,10 +339,38 @@ func TestAllowedContextMappings_DropsOverLongAndDuplicateKeys(t *testing.T) { } } +func TestAllowedContextMappings_CountsKeyLengthInRunes(t *testing.T) { + kept := strings.Repeat("键", 40) + dropped := strings.Repeat("键", 65) + setAllowlist(t, kept+","+dropped) + + got := allowedContextMappings() + if len(got) != 1 || got[0].source != kept { + t.Fatalf("got %#v, want the 40-rune CJK key kept and the 65-rune key dropped", got) + } +} + +func TestAllowedContextMappings_CachesUntilReset(t *testing.T) { + setAllowlist(t, "a") + if got := allowedContextMappings(); len(got) != 1 || got[0].source != "a" { + t.Fatalf("got %#v, want [a]", got) + } + + t.Setenv(traceContextKeysEnvVar, "b") + if got := allowedContextMappings(); len(got) != 1 || got[0].source != "a" { + t.Errorf("OnceValue lost the first parse: %#v", got) + } + + resetAllowedContextMappings() + if got := allowedContextMappings(); len(got) != 1 || got[0].source != "b" { + t.Errorf("reset did not reparse: %#v", got) + } +} + // Langfuse and comparable backends filter on attributes present on each span, // so the promoted values must reach descendants, not just the root span. func TestCallerContextAttributes_ReachEverySpan(t *testing.T) { - t.Setenv(traceContextKeysEnvVar, `[{"from":"sub","to":"user.id"}]`) + setAllowlist(t, `[{"from":"sub","to":"user.id"}]`) exporter := tracetest.NewInMemoryExporter() tp := sdktrace.NewTracerProvider( @@ -317,7 +405,7 @@ func TestCallerContextAttributes_ReachEverySpan(t *testing.T) { } func TestCallerContextAttributes_DisabledLeavesSpansUnchanged(t *testing.T) { - t.Setenv(traceContextKeysEnvVar, "") + setAllowlist(t, "") exporter := tracetest.NewInMemoryExporter() tp := sdktrace.NewTracerProvider( @@ -345,7 +433,7 @@ func TestCallerContextAttributes_DisabledLeavesSpansUnchanged(t *testing.T) { // Custom keys still cannot shadow a semantic convention attribute such as // service.name. Registry names are the documented exception. func TestCallerContextAttributes_CannotShadowSemanticConventions(t *testing.T) { - t.Setenv(traceContextKeysEnvVar, "service.name") + setAllowlist(t, "service.name") got := CallerContextAttributes(context.Background(), map[string]any{"service.name": "impostor"}) diff --git a/go/core/pkg/env/otel.go b/go/core/pkg/env/otel.go index 2b8775b653..4145a729be 100644 --- a/go/core/pkg/env/otel.go +++ b/go/core/pkg/env/otel.go @@ -43,9 +43,9 @@ var ( "", "Allowlist of caller-supplied context keys promoted onto every agent span. "+ "Accepts a comma-separated list of source keys, or a JSON array of strings and "+ - "{from, to, hash} objects. Registry names (user.*, enduser.*, session.id) are left "+ - "unprefixed; everything else is emitted as kagent.context. unless the name is "+ - "already in the kagent. namespace. Empty (the default) disables promotion.", + "{from, to, hash} objects. Registry names (user.id, user.hash, enduser.id, session.id) "+ + "are left unprefixed; everything else is emitted as kagent.context.. "+ + "Empty (the default) disables promotion.", ComponentAgentRuntime, ) diff --git a/helm/kagent/tests/controller-deployment_test.yaml b/helm/kagent/tests/controller-deployment_test.yaml index b2f3ad4b54..26ad759c8c 100644 --- a/helm/kagent/tests/controller-deployment_test.yaml +++ b/helm/kagent/tests/controller-deployment_test.yaml @@ -857,12 +857,12 @@ tests: - from: sub to: user.id - from: thread_id - to: kagent.thread_id + to: session.id - channel asserts: - equal: path: data.KAGENT_TRACE_CONTEXT_KEYS - value: '[{"from":"sub","to":"user.id"},{"from":"thread_id","to":"kagent.thread_id"},"channel"]' + value: '[{"from":"sub","to":"user.id"},{"from":"thread_id","to":"session.id"},"channel"]' - it: should omit KAGENT_TRACE_CONTEXT_HASH_KEY by default template: controller-deployment.yaml diff --git a/helm/kagent/values.yaml b/helm/kagent/values.yaml index 31942c1300..092f4194a0 100644 --- a/helm/kagent/values.yaml +++ b/helm/kagent/values.yaml @@ -854,13 +854,12 @@ otel: # # Entries may be a source key or a mapping. Prefer an opaque identifier # such as an OIDC subject for user.id; do not put names or email addresses - # on spans. Registry names (user.*, enduser.*, session.id) are left - # unprefixed; names already in the kagent. namespace are left as-is; - # everything else is emitted as kagent.context.. + # on spans. Registry names (user.id, user.hash, enduser.id, session.id) + # are left unprefixed; everything else is emitted as kagent.context.. # # contextKeys: # - {from: sub, to: user.id} - # - {from: thread_id, to: kagent.thread_id} + # - thread_id # - channel # # To derive a stable identifier without putting the original value on the diff --git a/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py b/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py index fb8cf98c5f..9ebe243eaf 100644 --- a/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py +++ b/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py @@ -33,9 +33,8 @@ get_kagent_metadata_key, hitl_activated, now_timestamp, - read_message_metadata, ) -from kagent.core.tracing import caller_context_attributes +from kagent.core.tracing import merge_caller_context_attributes from kagent.core.tracing._span_processor import clear_kagent_span_attributes, set_kagent_span_attributes from pydantic import BaseModel @@ -135,8 +134,8 @@ async def execute(self, context: RequestContext, event_queue: EventQueue) -> Non } # Allowlisted caller context joins the request-scoped bag rather # than a single span, so tool, sub-agent, and model spans all - # carry it. - span_attributes.update(caller_context_attributes(read_message_metadata(context.message))) + # carry it. Fill-if-absent so a caller cannot override kagent.user_id. + merge_caller_context_attributes(span_attributes, message=context.message) context_token = set_kagent_span_attributes( {key: value for key, value in span_attributes.items() if value is not None} ) diff --git a/python/packages/kagent-core/src/kagent/core/tracing/__init__.py b/python/packages/kagent-core/src/kagent/core/tracing/__init__.py index 87307f1e71..96d6cdcd40 100644 --- a/python/packages/kagent-core/src/kagent/core/tracing/__init__.py +++ b/python/packages/kagent-core/src/kagent/core/tracing/__init__.py @@ -1,4 +1,4 @@ -from ._context_attributes import caller_context_attributes +from ._context_attributes import caller_context_attributes, merge_caller_context_attributes from ._utils import configure, force_flush -__all__ = ["caller_context_attributes", "configure", "force_flush"] +__all__ = ["caller_context_attributes", "configure", "force_flush", "merge_caller_context_attributes"] diff --git a/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py b/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py index 67853b3945..22ef891a4a 100644 --- a/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py +++ b/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py @@ -2,6 +2,7 @@ from __future__ import annotations +import functools import hashlib import hmac import json @@ -9,9 +10,12 @@ from dataclasses import dataclass from typing import Any, Optional +from a2a.types import Message from opentelemetry import baggage from opentelemetry import context as otel_context +from kagent.core.a2a._consts import read_message_metadata + # Allowlist of caller-supplied context keys to promote onto agent spans. # Unset or empty (the default) disables promotion entirely. # @@ -25,9 +29,9 @@ TRACE_CONTEXT_HASH_KEY_ENV_VAR = "KAGENT_TRACE_CONTEXT_HASH_KEY" # Namespaces custom promoted values so they cannot shadow a semantic -# convention attribute such as ``service.name``. Registry names -# (``user.*``, ``enduser.*``, ``session.id``) and names already in the -# ``kagent.`` namespace are left unprefixed; see ``_span_attribute_name``. +# convention attribute such as ``service.name``. Only the explicit registry +# names in ``_UNPREFIXED_ATTRIBUTES`` pass through as-is; see +# ``_span_attribute_name``. CONTEXT_ATTRIBUTE_PREFIX = "kagent.context." HASH_HMAC_SHA256 = "hmac-sha256" @@ -36,6 +40,11 @@ MAX_CONTEXT_KEY_LENGTH = 64 MAX_CONTEXT_VALUE_LENGTH = 256 +# OpenTelemetry registry names callers may emit without a kagent.context. +# prefix. An invented name such as user.asdasd cannot masquerade as a +# semantic convention. +_UNPREFIXED_ATTRIBUTES = frozenset({"user.id", "user.hash", "enduser.id", "session.id"}) + @dataclass(frozen=True) class _ContextMapping: @@ -47,6 +56,7 @@ class _ContextMapping: def caller_context_attributes( metadata: Optional[dict[str, Any]] = None, context: Optional[otel_context.Context] = None, + message: Optional[Message] = None, ) -> dict[str, str]: """Return the caller context values an operator allowlisted for tracing. @@ -57,13 +67,20 @@ def caller_context_attributes( Values are read from W3C baggage first and then from the A2A message metadata, which is the more specific source for a single message and - therefore wins. Both are untrusted input, so keys must appear in the - allowlist, values are stripped of control characters and truncated, and - attribute names go through ``_span_attribute_name``. + therefore wins — but only when the sanitised metadata scalar is non-empty, + so a blank metadata entry cannot wipe a baggage value. Both are untrusted + input, so keys must appear in the allowlist, values are stripped of control + characters and truncated, and attribute names go through + ``_span_attribute_name``. + + ``message`` is decoded only after the allowlist is known to be non-empty, + so the default (promotion off) does not pay for protobuf Struct conversion. Args: metadata: A2A message metadata as a plain dict (may be ``None``). context: OTel context to read baggage from. Defaults to the current one. + message: A2A message whose metadata is decoded lazily when ``metadata`` + is ``None`` and the allowlist is non-empty. Returns: Attribute name to sanitised value. Empty when the allowlist is empty, @@ -72,6 +89,8 @@ def caller_context_attributes( mappings = _allowed_context_mappings() if not mappings: return {} + if metadata is None and message is not None: + metadata = read_message_metadata(message) bag = baggage.get_all(context) attributes: dict[str, str] = {} @@ -80,7 +99,9 @@ def caller_context_attributes( if metadata is not None: scalar = _scalar_string(metadata.get(mapping.source)) if scalar is not None: - value = _sanitize_context_value(scalar) + cleaned = _sanitize_context_value(scalar) + if cleaned: + value = cleaned if not value: continue if mapping.hash: @@ -96,19 +117,37 @@ def caller_context_attributes( return attributes -def _allowed_context_mappings() -> list[_ContextMapping]: +def merge_caller_context_attributes( + existing: dict[str, Any], + metadata: Optional[dict[str, Any]] = None, + context: Optional[otel_context.Context] = None, + message: Optional[Message] = None, +) -> None: + """Copy promoted caller context into *existing* without replacing keys. + + Caller-supplied data cannot override ``kagent.user_id``, ``kagent.app_name``, + or other attributes the runtime already stamped. + """ + for key, value in caller_context_attributes(metadata, context, message=message).items(): + if key not in existing: + existing[key] = value + + +@functools.cache +def _allowed_context_mappings() -> tuple[_ContextMapping, ...]: """Parse the ``KAGENT_TRACE_CONTEXT_KEYS`` allowlist. Keys that are empty, over-long, or contain whitespace or control characters are dropped, and the list is capped at ``MAX_CONTEXT_KEYS`` so a misconfigured allowlist cannot inflate span cardinality without bound. + Cached because the allowlist cannot change without a process restart. """ raw = os.getenv(TRACE_CONTEXT_KEYS_ENV_VAR, "").strip() if not raw: - return [] + return () if raw.startswith("["): - return _cap_mappings(_parse_json_allowlist(raw)) - return _cap_mappings(_parse_comma_allowlist(raw)) + return tuple(_cap_mappings(_parse_json_allowlist(raw))) + return tuple(_cap_mappings(_parse_comma_allowlist(raw))) def _parse_comma_allowlist(raw: str) -> list[_ContextMapping]: @@ -180,21 +219,16 @@ def _cap_mappings(mappings: list[_ContextMapping]) -> list[_ContextMapping]: def _span_attribute_name(name: str) -> str: """Return the name written onto the span. - ``user.*``, ``enduser.*``, and ``session.id`` pass through unprefixed so - operators can use the semantic convention names. Names already in the - ``kagent.`` namespace are left as-is. Everything else is placed under - ``kagent.context.`` so a caller-supplied ``service.name`` cannot shadow - the real one. + ``user.id``, ``user.hash``, ``enduser.id``, and ``session.id`` pass through + unprefixed so operators can use the semantic convention names. Everything + else is placed under ``kagent.context.`` so a caller-supplied + ``service.name`` or ``kagent.user_id`` cannot shadow the real one. """ - if _is_registry_attribute(name) or name.startswith("kagent."): + if name in _UNPREFIXED_ATTRIBUTES: return name return CONTEXT_ATTRIBUTE_PREFIX + name -def _is_registry_attribute(name: str) -> bool: - return name.startswith("user.") or name.startswith("enduser.") or name == "session.id" - - def _is_attribute_key(key: str) -> bool: """Report whether *key* is safe to use as a span attribute name.""" return not any(_is_control(char) or char.isspace() for char in key) @@ -218,14 +252,20 @@ def _scalar_string(value: Any) -> Optional[str]: if isinstance(value, int): return str(value) if isinstance(value, float): - # protobuf Struct has a single numeric type, so a JSON integer arrives - # as a float. Render it without the trailing ".0" to match the Go ADK. - return str(int(value)) if value.is_integer() else repr(value) + return _format_float(value) if isinstance(value, str): return value return None +def _format_float(value: float) -> str: + """Match Go ``strconv.FormatFloat(v, 'f', -1, 64)``: never scientific notation.""" + formatted = format(value, ".16f") + if "." in formatted: + formatted = formatted.rstrip("0").rstrip(".") + return formatted or "0" + + def _sanitize_context_value(value: Any) -> str: """Drop control characters and trim space so a value cannot forge structure.""" if not isinstance(value, str): diff --git a/python/packages/kagent-core/tests/conftest.py b/python/packages/kagent-core/tests/conftest.py new file mode 100644 index 0000000000..1868bcf802 --- /dev/null +++ b/python/packages/kagent-core/tests/conftest.py @@ -0,0 +1,11 @@ +import pytest + +from kagent.core.tracing._context_attributes import _allowed_context_mappings + + +@pytest.fixture(autouse=True) +def clear_allowed_context_mappings_cache(): + """The allowlist is cached for process lifetime; tests change the env var.""" + _allowed_context_mappings.cache_clear() + yield + _allowed_context_mappings.cache_clear() diff --git a/python/packages/kagent-core/tests/test_caller_context_attributes.py b/python/packages/kagent-core/tests/test_caller_context_attributes.py index 2584f84ccf..e942ba4f33 100644 --- a/python/packages/kagent-core/tests/test_caller_context_attributes.py +++ b/python/packages/kagent-core/tests/test_caller_context_attributes.py @@ -1,5 +1,6 @@ import hashlib import hmac +from unittest.mock import MagicMock, patch import pytest from opentelemetry import baggage @@ -8,7 +9,7 @@ from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter -from kagent.core.tracing import caller_context_attributes +from kagent.core.tracing import caller_context_attributes, merge_caller_context_attributes from kagent.core.tracing._context_attributes import ( CONTEXT_ATTRIBUTE_PREFIX, MAX_CONTEXT_KEY_LENGTH, @@ -70,6 +71,13 @@ def test_message_metadata_overrides_baggage(self, monkeypatch): baggage_context({"sub": "from-baggage"}), ) == {"kagent.context.sub": "from-metadata"} + def test_empty_metadata_does_not_override_baggage(self, monkeypatch): + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "sub") + assert caller_context_attributes( + {"sub": " "}, + baggage_context({"sub": "from-baggage"}), + ) == {"kagent.context.sub": "from-baggage"} + def test_ignores_keys_outside_the_allowlist(self, monkeypatch): monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "thread_id") assert caller_context_attributes( @@ -86,6 +94,7 @@ def test_ignores_keys_outside_the_allowlist(self, monkeypatch): # protobuf Struct has one numeric type, so JSON integers arrive as floats. (3.0, "3"), (2.5, "2.5"), + (1e20, "100000000000000000000"), ], ) def test_renders_scalar_metadata_types(self, monkeypatch, value, expected): @@ -116,22 +125,34 @@ def test_cannot_shadow_semantic_conventions(self, monkeypatch): assert promoted == {f"{CONTEXT_ATTRIBUTE_PREFIX}service.name": "impostor"} def test_registry_attributes_stay_unprefixed(self, monkeypatch): - """user.*, enduser.*, and session.id are semantic convention names.""" - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "user.id,enduser.id,session.id,channel") + """Only the explicit registry set is left unprefixed.""" + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "user.id,user.hash,enduser.id,session.id,channel") assert caller_context_attributes( { "user.id": "opaque-subject", + "user.hash": "abc", "enduser.id": "end-user", "session.id": "sess-1", "channel": "C0AB1", } ) == { "user.id": "opaque-subject", + "user.hash": "abc", "enduser.id": "end-user", "session.id": "sess-1", "kagent.context.channel": "C0AB1", } + def test_unknown_user_attributes_are_prefixed(self, monkeypatch): + monkeypatch.setenv( + TRACE_CONTEXT_KEYS_ENV_VAR, + '[{"from":"x","to":"user.asdasd"},{"from":"y","to":"user.email"}]', + ) + assert caller_context_attributes({"x": "nope", "y": "also-nope"}) == { + "kagent.context.user.asdasd": "nope", + "kagent.context.user.email": "also-nope", + } + def test_session_id_is_unprefixed_but_session_foo_is_not(self, monkeypatch): monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "session.id,session.foo") assert caller_context_attributes({"session.id": "sess-1", "session.foo": "other"}) == { @@ -139,17 +160,27 @@ def test_session_id_is_unprefixed_but_session_foo_is_not(self, monkeypatch): "kagent.context.session.foo": "other", } - def test_maps_source_keys_onto_registry_and_kagent_names(self, monkeypatch): + def test_maps_source_keys_onto_registry_names(self, monkeypatch): monkeypatch.setenv( TRACE_CONTEXT_KEYS_ENV_VAR, - '[{"from":"sub","to":"user.id"},{"from":"thread_id","to":"kagent.thread_id"},"channel"]', + '[{"from":"sub","to":"user.id"},{"from":"thread_id","to":"session.id"},"channel"]', ) assert caller_context_attributes({"sub": "opaque-subject", "thread_id": "T123", "channel": "C0AB1"}) == { "user.id": "opaque-subject", - "kagent.thread_id": "T123", + "session.id": "T123", "kagent.context.channel": "C0AB1", } + def test_kagent_destination_names_are_prefixed(self, monkeypatch): + monkeypatch.setenv( + TRACE_CONTEXT_KEYS_ENV_VAR, + '[{"from":"uid","to":"kagent.user_id"},{"from":"tid","to":"kagent.thread_id"}]', + ) + assert caller_context_attributes({"uid": "attacker", "tid": "T123"}) == { + "kagent.context.kagent.user_id": "attacker", + "kagent.context.kagent.thread_id": "T123", + } + def test_invalid_json_allowlist_promotes_nothing(self, monkeypatch): monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, '[{"from":"sub"') assert caller_context_attributes({"sub": "opaque-subject"}) == {} @@ -182,6 +213,43 @@ def test_unknown_hash_emits_nothing(self, monkeypatch): monkeypatch.setenv(TRACE_CONTEXT_HASH_KEY_ENV_VAR, "test-hmac-key") assert caller_context_attributes({"email": "ada@example.com"}) == {} + def test_skips_message_decode_when_allowlist_empty(self, monkeypatch): + monkeypatch.delenv(TRACE_CONTEXT_KEYS_ENV_VAR, raising=False) + with patch( + "kagent.core.tracing._context_attributes.read_message_metadata", + return_value={"thread_id": "T1"}, + ) as decode: + assert caller_context_attributes(message=MagicMock()) == {} + decode.assert_not_called() + + def test_decodes_message_when_allowlist_is_set(self, monkeypatch): + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "thread_id") + message = MagicMock() + with patch( + "kagent.core.tracing._context_attributes.read_message_metadata", + return_value={"thread_id": "T1"}, + ) as decode: + assert caller_context_attributes(message=message) == {"kagent.context.thread_id": "T1"} + decode.assert_called_once_with(message) + + +class TestMergeCallerContextAttributes: + def test_does_not_override_existing_keys(self, monkeypatch): + monkeypatch.setenv( + TRACE_CONTEXT_KEYS_ENV_VAR, + '[{"from":"sub","to":"user.id"},"thread_id"]', + ) + existing = {"user.id": "runtime-user", "kagent.user_id": "runtime-user"} + merge_caller_context_attributes( + existing, + {"sub": "attacker", "thread_id": "T1"}, + ) + assert existing == { + "user.id": "runtime-user", + "kagent.user_id": "runtime-user", + "kagent.context.thread_id": "T1", + } + class TestAllowedContextMappings: """Tests for allowlist parsing and its bounds.""" @@ -195,13 +263,27 @@ def test_drops_over_long_and_duplicate_keys(self, monkeypatch): monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, f"a,a,{'b' * (MAX_CONTEXT_KEY_LENGTH + 1)},c") assert [mapping.source for mapping in _allowed_context_mappings()] == ["a", "c"] + def test_counts_key_length_in_code_points(self, monkeypatch): + kept = "键" * 40 + dropped = "键" * 65 + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, f"{kept},{dropped}") + assert [mapping.source for mapping in _allowed_context_mappings()] == [kept] + def test_drops_keys_that_are_not_valid_attribute_names(self, monkeypatch): monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "good, bad key ,\tanother\tbad") assert [mapping.source for mapping in _allowed_context_mappings()] == ["good"] def test_empty_allowlist_disables_promotion(self, monkeypatch): monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, " , ,") - assert _allowed_context_mappings() == [] + assert _allowed_context_mappings() == () + + def test_cache_survives_env_change_until_cleared(self, monkeypatch): + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "a") + assert [mapping.source for mapping in _allowed_context_mappings()] == ["a"] + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "b") + assert [mapping.source for mapping in _allowed_context_mappings()] == ["a"] + _allowed_context_mappings.cache_clear() + assert [mapping.source for mapping in _allowed_context_mappings()] == ["b"] class TestPromotedAttributesReachEverySpan: diff --git a/python/packages/kagent-crewai/src/kagent/crewai/_executor.py b/python/packages/kagent-crewai/src/kagent/crewai/_executor.py index dc5a80b103..c99ff97e82 100644 --- a/python/packages/kagent-crewai/src/kagent/crewai/_executor.py +++ b/python/packages/kagent-crewai/src/kagent/crewai/_executor.py @@ -22,8 +22,8 @@ TaskStatusUpdateEvent, ) from google.protobuf.json_format import MessageToDict -from kagent.core.a2a import get_kagent_metadata_key, now_timestamp, read_message_metadata -from kagent.core.tracing import caller_context_attributes +from kagent.core.a2a import get_kagent_metadata_key, now_timestamp +from kagent.core.tracing import merge_caller_context_attributes from kagent.core.tracing._span_processor import ( clear_kagent_span_attributes, set_kagent_span_attributes, @@ -196,6 +196,7 @@ def _convert_a2a_request_to_span_attributes( # Allowlisted caller context joins the request-scoped bag rather than a # single span, so tool, sub-agent, and model spans all carry it. - span_attributes.update(caller_context_attributes(read_message_metadata(request.message))) + # Fill-if-absent so a caller cannot override kagent.user_id. + merge_caller_context_attributes(span_attributes, message=request.message) return span_attributes diff --git a/python/packages/kagent-langgraph/src/kagent/langgraph/_executor.py b/python/packages/kagent-langgraph/src/kagent/langgraph/_executor.py index 46c76ddc3a..76f94aa42e 100644 --- a/python/packages/kagent-langgraph/src/kagent/langgraph/_executor.py +++ b/python/packages/kagent-langgraph/src/kagent/langgraph/_executor.py @@ -40,11 +40,10 @@ get_tool_approval_response, hitl_activated, now_timestamp, - read_message_metadata, require_ask_user_response, require_tool_approval_response, ) -from kagent.core.tracing import caller_context_attributes +from kagent.core.tracing import merge_caller_context_attributes from kagent.core.tracing._span_processor import ( clear_kagent_span_attributes, set_kagent_span_attributes, @@ -545,6 +544,7 @@ def _convert_a2a_request_to_span_attributes( # Allowlisted caller context joins the request-scoped bag rather than a # single span, so tool, sub-agent, and model spans all carry it. - span_attributes.update(caller_context_attributes(read_message_metadata(request.message))) + # Fill-if-absent so a caller cannot override kagent.user_id. + merge_caller_context_attributes(span_attributes, message=request.message) return span_attributes From 49003ccb583af7460fb47a487765372f97d2100d Mon Sep 17 00:00:00 2001 From: Ricardo Temperini <29879569+rtemperini@users.noreply.github.com> Date: Mon, 31 Aug 2026 11:00:32 +0200 Subject: [PATCH 07/13] fix: fail closed on malformed hash and match tiny floats Reject non-string hash/to so Python cannot emit plaintext, format 1e-20 as Go does, and stamp caller context from the OpenAI executor. Signed-off-by: Ricardo Temperini <29879569+rtemperini@users.noreply.github.com> --- docs/architecture/trace-context.md | 6 +- .../pkg/telemetry/context_attributes_test.go | 19 +++++ .../tests/unittests/test_agent_executor.py | 45 ++++++++++- .../core/tracing/_context_attributes.py | 75 ++++++++++++++++--- .../tests/test_caller_context_attributes.py | 19 +++++ .../kagent-crewai/tests/test_executor.py | 35 ++++++++- .../src/kagent/openai/_agent_executor.py | 45 +++++++++++ .../tests/test_caller_context.py | 49 ++++++++++++ 8 files changed, 278 insertions(+), 15 deletions(-) create mode 100644 python/packages/kagent-openai/tests/test_caller_context.py diff --git a/docs/architecture/trace-context.md b/docs/architecture/trace-context.md index 10c45c8b54..af73c4a6e0 100644 --- a/docs/architecture/trace-context.md +++ b/docs/architecture/trace-context.md @@ -173,7 +173,11 @@ processors: | Component | Path | |---|---| | Go ADK | `go/adk/pkg/telemetry/context_attributes.go` | -| Python | `python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py` | +| Python core | `python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py` | +| Python ADK | `python/packages/kagent-adk/src/kagent/adk/_agent_executor.py` | +| Python LangGraph | `python/packages/kagent-langgraph/src/kagent/langgraph/_executor.py` | +| Python CrewAI | `python/packages/kagent-crewai/src/kagent/crewai/_executor.py` | +| Python OpenAI | `python/packages/kagent-openai/src/kagent/openai/_agent_executor.py` | | Controller forwarding | `go/core/v2/translator/kagent/compiler.go` | | Helm | `helm/kagent/templates/controller-configmap.yaml` | diff --git a/go/adk/pkg/telemetry/context_attributes_test.go b/go/adk/pkg/telemetry/context_attributes_test.go index c3cff09b7c..3200c711f4 100644 --- a/go/adk/pkg/telemetry/context_attributes_test.go +++ b/go/adk/pkg/telemetry/context_attributes_test.go @@ -121,6 +121,12 @@ func TestCallerContextAttributes(t *testing.T) { metadata: map[string]any{"value": float64(1e20)}, want: map[string]string{"kagent.context.value": "100000000000000000000"}, }, + { + name: "formats tiny floats without rounding to zero", + allowlist: "value", + metadata: map[string]any{"value": float64(1e-20)}, + want: map[string]string{"kagent.context.value": "0.00000000000000000001"}, + }, { name: "skips non-scalar and empty metadata values", allowlist: "nested,list,blank", @@ -208,6 +214,19 @@ func TestCallerContextAttributes(t *testing.T) { metadata: map[string]any{"sub": "opaque-subject"}, want: nil, }, + { + name: "non-string hash does not emit plaintext", + allowlist: `[{"from":"email","to":"user.hash","hash":123}]`, + hashKey: "test-hmac-key", + metadata: map[string]any{"email": "ada@example.com"}, + want: nil, + }, + { + name: "non-string to drops the mapping", + allowlist: `[{"from":"thread_id","to":123}]`, + metadata: map[string]any{"thread_id": "T1"}, + want: nil, + }, } for _, tt := range tests { diff --git a/python/packages/kagent-adk/tests/unittests/test_agent_executor.py b/python/packages/kagent-adk/tests/unittests/test_agent_executor.py index 15215ee949..393cf5f615 100644 --- a/python/packages/kagent-adk/tests/unittests/test_agent_executor.py +++ b/python/packages/kagent-adk/tests/unittests/test_agent_executor.py @@ -1,6 +1,6 @@ from __future__ import annotations -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock import pytest from a2a.server.agent_execution.context import RequestContext @@ -104,3 +104,46 @@ async def execute(self, request_context, queue): assert calls["config"].request_converter == executor._convert_request executor._prepare_session.assert_awaited_once_with(context, run_request, runner) executor._safe_close_runner.assert_awaited_once_with(runner) + + +@pytest.mark.asyncio +async def test_execute_merges_caller_context_without_eager_decode(monkeypatch): + context = _request_context() + event_queue = object() + runner = object() + run_request = AgentRunRequest( + user_id="user-1", + session_id="context-1", + run_config=RunConfig(), + ) + executor = A2aAgentExecutor(runner=lambda: None) + executor._resolve_runner = AsyncMock(return_value=runner) + executor._convert_request = lambda request_context, part_converter: run_request + executor._prepare_session = AsyncMock() + executor._safe_close_runner = AsyncMock() + + recorded = {} + + def fake_merge(existing, metadata=None, context=None, message=None): + recorded["existing"] = dict(existing) + recorded["message"] = message + + monkeypatch.setattr(executor_module, "merge_caller_context_attributes", fake_merge) + decode = MagicMock(return_value={"thread_id": "T1"}) + if hasattr(executor_module, "read_message_metadata"): + monkeypatch.setattr(executor_module, "read_message_metadata", decode) + + class FakeUpstreamExecutor: + def __init__(self, *, runner, config, force_new_version): + pass + + async def execute(self, request_context, queue): + return None + + monkeypatch.setattr(executor_module, "UpstreamA2aAgentExecutor", FakeUpstreamExecutor) + + await executor.execute(context, event_queue) + + assert recorded["message"] is context.message + assert recorded["existing"]["kagent.user_id"] == "user-1" + decode.assert_not_called() diff --git a/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py b/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py index 22ef891a4a..ac90f8d2a9 100644 --- a/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py +++ b/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py @@ -171,15 +171,15 @@ def _parse_json_allowlist(raw: str) -> list[_ContextMapping]: if isinstance(item, str): mapping = _new_context_mapping(item, "", "") elif isinstance(item, dict): - from_key, to_key, hash_alg = item.get("from"), item.get("to"), item.get("hash") - if from_key is not None and not isinstance(from_key, str): + from_key, from_ok = _json_string_field(item, "from") + to_key, to_ok = _json_string_field(item, "to") + hash_alg, hash_ok = _json_string_field(item, "hash") + if not from_ok or not to_ok or not hash_ok: + # Non-string hash/to must drop the mapping. Coercing hash: 123 + # to "no hash" would emit the original value in plaintext. mapping = None else: - mapping = _new_context_mapping( - from_key or "", - to_key if isinstance(to_key, str) else "", - hash_alg if isinstance(hash_alg, str) else "", - ) + mapping = _new_context_mapping(from_key, to_key, hash_alg) else: mapping = None if mapping is not None: @@ -187,6 +187,23 @@ def _parse_json_allowlist(raw: str) -> list[_ContextMapping]: return mappings +def _json_string_field(item: dict[str, Any], name: str) -> tuple[str, bool]: + """Return a JSON string field, or reject it when the type is invalid. + + Missing keys and JSON null are unset (empty string), matching Go's + encoding/json behaviour for string fields. Any other non-string type is + a malformed mapping and must fail closed. + """ + if name not in item: + return "", True + raw = item[name] + if raw is None: + return "", True + if isinstance(raw, str): + return raw, True + return "", False + + def _new_context_mapping(from_key: str, to_key: str, hash_alg: str) -> Optional[_ContextMapping]: from_key = from_key.strip() to_key = to_key.strip() @@ -259,11 +276,45 @@ def _scalar_string(value: Any) -> Optional[str]: def _format_float(value: float) -> str: - """Match Go ``strconv.FormatFloat(v, 'f', -1, 64)``: never scientific notation.""" - formatted = format(value, ".16f") - if "." in formatted: - formatted = formatted.rstrip("0").rstrip(".") - return formatted or "0" + """Match Go ``strconv.FormatFloat(v, 'f', -1, 64)``: never scientific notation. + + ``format(value, '.16f')`` rounds values such as ``1e-20`` to ``0``. Expand + Python's shortest unique repr from scientific form into fixed-point + instead, so tiny values keep their digits. + """ + return _fixed_point_from_shortest(repr(value)) + + +def _fixed_point_from_shortest(text: str) -> str: + lower = text.lower() + negative = lower.startswith("-") + if negative: + lower = lower[1:] + if "e" not in lower: + if "." in lower: + lower = lower.rstrip("0").rstrip(".") + out = lower or "0" + return f"-{out}" if negative and out != "0" else out + mantissa, exp_s = lower.split("e") + exp = int(exp_s) + if "." in mantissa: + whole, frac = mantissa.split(".") + digits = whole + frac + exp -= len(frac) + else: + digits = mantissa + digits = digits.lstrip("0") or "0" + if exp >= 0: + out = digits + "0" * exp + else: + shift = -exp + if len(digits) <= shift: + out = "0." + "0" * (shift - len(digits)) + digits + else: + idx = len(digits) - shift + frac = digits[idx:].rstrip("0") + out = digits[:idx] + (("." + frac) if frac else "") + return f"-{out}" if negative and out != "0" else out def _sanitize_context_value(value: Any) -> str: diff --git a/python/packages/kagent-core/tests/test_caller_context_attributes.py b/python/packages/kagent-core/tests/test_caller_context_attributes.py index e942ba4f33..25ec635d67 100644 --- a/python/packages/kagent-core/tests/test_caller_context_attributes.py +++ b/python/packages/kagent-core/tests/test_caller_context_attributes.py @@ -95,6 +95,7 @@ def test_ignores_keys_outside_the_allowlist(self, monkeypatch): (3.0, "3"), (2.5, "2.5"), (1e20, "100000000000000000000"), + (1e-20, "0.00000000000000000001"), ], ) def test_renders_scalar_metadata_types(self, monkeypatch, value, expected): @@ -213,6 +214,24 @@ def test_unknown_hash_emits_nothing(self, monkeypatch): monkeypatch.setenv(TRACE_CONTEXT_HASH_KEY_ENV_VAR, "test-hmac-key") assert caller_context_attributes({"email": "ada@example.com"}) == {} + def test_non_string_hash_does_not_emit_plaintext(self, monkeypatch): + """A numeric hash field must drop the mapping, not skip hashing.""" + monkeypatch.setenv( + TRACE_CONTEXT_KEYS_ENV_VAR, + '[{"from":"email","to":"user.hash","hash":123}]', + ) + monkeypatch.setenv(TRACE_CONTEXT_HASH_KEY_ENV_VAR, "test-hmac-key") + promoted = caller_context_attributes({"email": "ada@example.com"}) + assert promoted == {} + assert not any("ada@example.com" in value for value in promoted.values()) + + def test_non_string_to_drops_the_mapping(self, monkeypatch): + monkeypatch.setenv( + TRACE_CONTEXT_KEYS_ENV_VAR, + '[{"from":"thread_id","to":123}]', + ) + assert caller_context_attributes({"thread_id": "T1"}) == {} + def test_skips_message_decode_when_allowlist_empty(self, monkeypatch): monkeypatch.delenv(TRACE_CONTEXT_KEYS_ENV_VAR, raising=False) with patch( diff --git a/python/packages/kagent-crewai/tests/test_executor.py b/python/packages/kagent-crewai/tests/test_executor.py index 5a04f491ba..a59751fea9 100644 --- a/python/packages/kagent-crewai/tests/test_executor.py +++ b/python/packages/kagent-crewai/tests/test_executor.py @@ -6,7 +6,7 @@ from google.protobuf.json_format import ParseDict from google.protobuf.struct_pb2 import Value -from kagent.crewai._executor import CrewAIAgentExecutor +from kagent.crewai._executor import CrewAIAgentExecutor, _convert_a2a_request_to_span_attributes def _request_context(*parts: Part) -> RequestContext: @@ -72,3 +72,36 @@ async def test_execute_falls_back_to_text_input_without_datapart(): crew.kickoff_async.assert_awaited_once_with(inputs={"input": "hello"}) _assert_content_artifact_closes_stream(events) + + +def test_span_attributes_skip_decode_when_allowlist_empty(monkeypatch): + from kagent.core.tracing._context_attributes import TRACE_CONTEXT_KEYS_ENV_VAR, _allowed_context_mappings + + monkeypatch.delenv(TRACE_CONTEXT_KEYS_ENV_VAR, raising=False) + _allowed_context_mappings.cache_clear() + decode = MagicMock(return_value={"thread_id": "T1"}) + monkeypatch.setattr("kagent.core.tracing._context_attributes.read_message_metadata", decode) + import kagent.crewai._executor as executor_module + + if hasattr(executor_module, "read_message_metadata"): + monkeypatch.setattr(executor_module, "read_message_metadata", decode) + + attrs = _convert_a2a_request_to_span_attributes(_request_context(Part(text="hello"))) + + decode.assert_not_called() + assert "kagent.context.thread_id" not in attrs + assert attrs["kagent.user_id"] + + +def test_span_attributes_do_not_override_runtime_user_id(monkeypatch): + fake = MagicMock(return_value={"kagent.user_id": "attacker", "kagent.context.thread_id": "T1"}) + monkeypatch.setattr("kagent.core.tracing._context_attributes.caller_context_attributes", fake) + import kagent.crewai._executor as executor_module + + if hasattr(executor_module, "caller_context_attributes"): + monkeypatch.setattr(executor_module, "caller_context_attributes", fake) + + attrs = _convert_a2a_request_to_span_attributes(_request_context(Part(text="hello"))) + + assert attrs["kagent.user_id"] != "attacker" + assert attrs["kagent.context.thread_id"] == "T1" diff --git a/python/packages/kagent-openai/src/kagent/openai/_agent_executor.py b/python/packages/kagent-openai/src/kagent/openai/_agent_executor.py index 62d1f5115e..b8c9321ce2 100644 --- a/python/packages/kagent-openai/src/kagent/openai/_agent_executor.py +++ b/python/packages/kagent-openai/src/kagent/openai/_agent_executor.py @@ -11,6 +11,7 @@ import uuid from collections.abc import Callable from dataclasses import dataclass +from typing import Any try: from typing import override # Python 3.12+ @@ -35,6 +36,11 @@ from agents.memory.session import SessionABC from agents.run import Runner from kagent.core.a2a import get_kagent_metadata_key, now_timestamp +from kagent.core.tracing import merge_caller_context_attributes +from kagent.core.tracing._span_processor import ( + clear_kagent_span_attributes, + set_kagent_span_attributes, +) from pydantic import BaseModel from ._event_converter import convert_openai_event_to_a2a_events @@ -173,6 +179,19 @@ async def execute( if not context.message: raise ValueError("A2A request must have a message") + span_attributes = _convert_a2a_request_to_span_attributes(context) + context_token = set_kagent_span_attributes(span_attributes) + try: + await self._execute_agent(context, event_queue) + finally: + clear_kagent_span_attributes(context_token) + + async def _execute_agent( + self, + context: RequestContext, + event_queue: EventQueue, + ) -> None: + # For new tasks, the first event must be a Task (not TaskStatusUpdateEvent). if not context.current_task: await event_queue.enqueue_event( @@ -277,3 +296,29 @@ async def execute( }, ) ) + + +def _get_user_id(request: RequestContext) -> str: + return getattr(request, "user_id", "admin@kagent.dev") + + +def _convert_a2a_request_to_span_attributes( + request: RequestContext, +) -> dict[str, Any]: + if not request.message: + raise ValueError("Request message cannot be None") + + span_attributes: dict[str, Any] = { + "kagent.user_id": _get_user_id(request), + "gen_ai.conversation.id": request.context_id, + } + + if request.task_id: + span_attributes["gen_ai.task.id"] = request.task_id + + # Allowlisted caller context joins the request-scoped bag rather than a + # single span, so tool, sub-agent, and model spans all carry it. + # Fill-if-absent so a caller cannot override kagent.user_id. + merge_caller_context_attributes(span_attributes, message=request.message) + + return span_attributes diff --git a/python/packages/kagent-openai/tests/test_caller_context.py b/python/packages/kagent-openai/tests/test_caller_context.py new file mode 100644 index 0000000000..5254523f71 --- /dev/null +++ b/python/packages/kagent-openai/tests/test_caller_context.py @@ -0,0 +1,49 @@ +from unittest.mock import MagicMock + +from a2a.server.agent_execution.context import RequestContext +from a2a.types import Message, Part, Role, SendMessageRequest + +from kagent.openai._agent_executor import _convert_a2a_request_to_span_attributes + + +def _request_context() -> RequestContext: + message = Message(role=Role.ROLE_USER, message_id="msg-1", parts=[Part(text="hello")]) + return RequestContext( + call_context=MagicMock(), + request=SendMessageRequest(message=message), + task_id="task-1", + context_id="ctx-1", + ) + + +def test_span_attributes_skip_decode_when_allowlist_empty(monkeypatch): + from kagent.core.tracing._context_attributes import TRACE_CONTEXT_KEYS_ENV_VAR, _allowed_context_mappings + + monkeypatch.delenv(TRACE_CONTEXT_KEYS_ENV_VAR, raising=False) + _allowed_context_mappings.cache_clear() + decode = MagicMock(return_value={"thread_id": "T1"}) + monkeypatch.setattr("kagent.core.tracing._context_attributes.read_message_metadata", decode) + import kagent.openai._agent_executor as executor_module + + if hasattr(executor_module, "read_message_metadata"): + monkeypatch.setattr(executor_module, "read_message_metadata", decode) + + attrs = _convert_a2a_request_to_span_attributes(_request_context()) + + decode.assert_not_called() + assert "kagent.context.thread_id" not in attrs + assert "kagent.user_id" in attrs + + +def test_span_attributes_do_not_override_runtime_user_id(monkeypatch): + fake = MagicMock(return_value={"kagent.user_id": "attacker", "kagent.context.thread_id": "T1"}) + monkeypatch.setattr("kagent.core.tracing._context_attributes.caller_context_attributes", fake) + import kagent.openai._agent_executor as executor_module + + if hasattr(executor_module, "caller_context_attributes"): + monkeypatch.setattr(executor_module, "caller_context_attributes", fake) + + attrs = _convert_a2a_request_to_span_attributes(_request_context()) + + assert attrs["kagent.user_id"] != "attacker" + assert attrs["kagent.context.thread_id"] == "T1" From cd9235afc303e14b188010ae9c79df904f86e3bc Mon Sep 17 00:00:00 2001 From: Ricardo Temperini <29879569+rtemperini@users.noreply.github.com> Date: Mon, 31 Aug 2026 11:06:53 +0200 Subject: [PATCH 08/13] fix: do not stamp allowlisted metadata as plaintext Skip KAGENT_TRACE_CONTEXT_KEYS source keys in generic A2A metadata span attributes so hashed or omitted values cannot leak on a2a.message.metadata.. Signed-off-by: Ricardo Temperini <29879569+rtemperini@users.noreply.github.com> --- go/adk/pkg/a2a/executor.go | 2 + go/adk/pkg/telemetry/attributes.go | 11 ++- go/adk/pkg/telemetry/attributes_test.go | 104 +++++++++++++++++++++ go/adk/pkg/telemetry/context_attributes.go | 75 ++++++++++++++- 4 files changed, 186 insertions(+), 6 deletions(-) diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index 10c48628ff..4c36a09701 100644 --- a/go/adk/pkg/a2a/executor.go +++ b/go/adk/pkg/a2a/executor.go @@ -137,6 +137,8 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.ExecutorCon ctx = telemetry.SetKAgentSpanAttributes(ctx, spanAttributes) ctx, invocationSpan := telemetry.StartInvocationSpan(ctx) defer invocationSpan.End() + // Allowlisted sources are skipped here so hashed or omitted values + // cannot leak as a2a.message.metadata. plaintext. telemetry.SetMessageMetadataAttributes(ctx, reqCtx.Message.Metadata) e.logger.Info("Execute", diff --git a/go/adk/pkg/telemetry/attributes.go b/go/adk/pkg/telemetry/attributes.go index ecfa64de74..9426fa7258 100644 --- a/go/adk/pkg/telemetry/attributes.go +++ b/go/adk/pkg/telemetry/attributes.go @@ -88,13 +88,22 @@ func stringAttributes(attrs map[string]string) []attribute.KeyValue { return out } -// SetMessageMetadataAttributes sets scalar values from an A2A message's metadata as span attributes. +// SetMessageMetadataAttributes sets scalar values from an A2A message's +// metadata as span attributes on the current span. +// +// Keys named in KAGENT_TRACE_CONTEXT_KEYS are skipped, including mappings +// that hash or omit the value. Stamping those sources here would put the +// original scalar on a2a.message.metadata. and bypass the allowlist. func SetMessageMetadataAttributes(ctx context.Context, metadata map[string]any) { if len(metadata) == 0 { return } + covered := policyCoveredContextSources() var attrs []attribute.KeyValue for k, v := range metadata { + if _, skip := covered[k]; skip { + continue + } key := "a2a.message.metadata." + k switch val := v.(type) { case string: diff --git a/go/adk/pkg/telemetry/attributes_test.go b/go/adk/pkg/telemetry/attributes_test.go index 4e689d68ec..4b26373e42 100644 --- a/go/adk/pkg/telemetry/attributes_test.go +++ b/go/adk/pkg/telemetry/attributes_test.go @@ -218,6 +218,110 @@ func TestSetMessageMetadataAttributes(t *testing.T) { } } +func TestSetMessageMetadataAttributes_DoesNotStampAllowlistedHashedPlaintext(t *testing.T) { + setAllowlist(t, `[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]`) + t.Setenv(traceContextHashKeyEnvVar, "test-hmac-key") + + attrs := recordMessageMetadata(t, map[string]any{ + "email": "ada@example.com", + "channel": "C0AB1", + }) + + if _, exists := attrs["a2a.message.metadata.email"]; exists { + t.Fatal("allowlisted hashed source must not be stamped as a2a.message.metadata.email") + } + assertNoPlaintextEmail(t, attrs) + if got := attrs["a2a.message.metadata.channel"].AsString(); got != "C0AB1" { + t.Errorf("non-allowlisted channel = %q, want C0AB1", got) + } +} + +func TestSetMessageMetadataAttributes_DoesNotStampWhenHMACKeyMissing(t *testing.T) { + setAllowlist(t, `[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]`) + t.Setenv(traceContextHashKeyEnvVar, "") + + attrs := recordMessageMetadata(t, map[string]any{"email": "ada@example.com"}) + + if _, exists := attrs["a2a.message.metadata.email"]; exists { + t.Fatal("omitted hashed source must not leak as a2a.message.metadata.email") + } + assertNoPlaintextEmail(t, attrs) +} + +func TestSetMessageMetadataAttributes_DoesNotStampMalformedHashSource(t *testing.T) { + setAllowlist(t, `[{"from":"email","to":"user.hash","hash":123}]`) + t.Setenv(traceContextHashKeyEnvVar, "test-hmac-key") + + attrs := recordMessageMetadata(t, map[string]any{"email": "ada@example.com"}) + + if _, exists := attrs["a2a.message.metadata.email"]; exists { + t.Fatal("malformed hash mapping must not leak plaintext as a2a.message.metadata.email") + } + assertNoPlaintextEmail(t, attrs) +} + +func TestMergeThenMessageMetadata_DoesNotLeakHashedPlaintext(t *testing.T) { + setAllowlist(t, `[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]`) + t.Setenv(traceContextHashKeyEnvVar, "test-hmac-key") + + exporter := tracetest.NewInMemoryExporter() + tp := sdktrace.NewTracerProvider( + sdktrace.WithSyncer(exporter), + sdktrace.WithSpanProcessor(kagentAttributesSpanProcessor{}), + ) + t.Cleanup(func() { _ = tp.Shutdown(context.Background()) }) + + metadata := map[string]any{"email": "ada@example.com"} + spanAttrs := map[string]string{} + MergeCallerContextAttributes(spanAttrs, context.Background(), metadata) + + tracer := tp.Tracer("test") + ctx, span := tracer.Start(context.Background(), "invocation") + ctx = SetKAgentSpanAttributes(ctx, spanAttrs) + SetMessageMetadataAttributes(ctx, metadata) + span.End() + + attrs := spanAttributesByName(t, exporter.GetSpans(), "invocation") + if _, exists := attrs["a2a.message.metadata.email"]; exists { + t.Fatal("hashed source leaked as a2a.message.metadata.email") + } + assertNoPlaintextEmail(t, attrs) + if got := attrs["user.hash"].AsString(); got == "" || strings.Contains(got, "@example.com") { + t.Errorf("user.hash = %q, want the HMAC digest", got) + } +} + +func recordMessageMetadata(t *testing.T, metadata map[string]any) map[string]attribute.Value { + t.Helper() + exporter := tracetest.NewInMemoryExporter() + tp := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter)) + t.Cleanup(func() { _ = tp.Shutdown(context.Background()) }) + + tracer := tp.Tracer("test") + ctx, span := tracer.Start(context.Background(), "test-span") + SetMessageMetadataAttributes(ctx, metadata) + span.End() + + spans := exporter.GetSpans() + if len(spans) == 0 { + t.Fatal("no spans recorded") + } + attrs := make(map[string]attribute.Value) + for _, a := range spans[0].Attributes { + attrs[string(a.Key)] = a.Value + } + return attrs +} + +func assertNoPlaintextEmail(t *testing.T, attrs map[string]attribute.Value) { + t.Helper() + for key, value := range attrs { + if strings.Contains(value.AsString(), "@example.com") { + t.Errorf("plaintext leaked onto %s = %q", key, value.AsString()) + } + } +} + func TestSetMessageMetadataAttributes_NilAndEmpty(t *testing.T) { exporter := tracetest.NewInMemoryExporter() tp := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter)) diff --git a/go/adk/pkg/telemetry/context_attributes.go b/go/adk/pkg/telemetry/context_attributes.go index 145c3e7689..e6c26e053a 100644 --- a/go/adk/pkg/telemetry/context_attributes.go +++ b/go/adk/pkg/telemetry/context_attributes.go @@ -55,17 +55,82 @@ type contextMapping struct { hash string } -// loadAllowedContextMappings parses KAGENT_TRACE_CONTEXT_KEYS once. The -// allowlist cannot change without a process restart; tests call +// contextPolicy is the parsed allowlist plus every source key named in it. +// Source keys are tracked even when a mapping is rejected (for example +// hash: 123), so generic metadata stamping cannot leak a value the operator +// asked to hash or omit. +type contextPolicy struct { + mappings []contextMapping + sources map[string]struct{} +} + +// loadContextPolicy parses KAGENT_TRACE_CONTEXT_KEYS once. The allowlist +// cannot change without a process restart; tests call // resetAllowedContextMappings after t.Setenv. -var loadAllowedContextMappings = sync.OnceValue(parseAllowedContextMappings) +var loadContextPolicy = sync.OnceValue(parseContextPolicy) func resetAllowedContextMappings() { - loadAllowedContextMappings = sync.OnceValue(parseAllowedContextMappings) + loadContextPolicy = sync.OnceValue(parseContextPolicy) } func allowedContextMappings() []contextMapping { - return loadAllowedContextMappings() + return loadContextPolicy().mappings +} + +func policyCoveredContextSources() map[string]struct{} { + return loadContextPolicy().sources +} + +func parseContextPolicy() contextPolicy { + return contextPolicy{ + mappings: parseAllowedContextMappings(), + sources: parseContextPolicySources(), + } +} + +func parseContextPolicySources() map[string]struct{} { + raw := strings.TrimSpace(os.Getenv(traceContextKeysEnvVar)) + if raw == "" { + return nil + } + if strings.HasPrefix(raw, "[") { + return jsonAllowlistSourceKeys(raw) + } + sources := make(map[string]struct{}) + for key := range strings.SplitSeq(raw, ",") { + if key = strings.TrimSpace(key); key != "" { + sources[key] = struct{}{} + } + } + return sources +} + +func jsonAllowlistSourceKeys(raw string) map[string]struct{} { + var items []json.RawMessage + if err := json.Unmarshal([]byte(raw), &items); err != nil { + return nil + } + sources := make(map[string]struct{}) + for _, item := range items { + var key string + if err := json.Unmarshal(item, &key); err == nil { + if key = strings.TrimSpace(key); key != "" { + sources[key] = struct{}{} + } + continue + } + // Only bind from so a non-string hash/to cannot hide the source key. + var spec struct { + From string `json:"from"` + } + if err := json.Unmarshal(item, &spec); err != nil { + continue + } + if key := strings.TrimSpace(spec.From); key != "" { + sources[key] = struct{}{} + } + } + return sources } // CallerContextAttributes returns the caller-supplied context values that an From 1de47c31e1b3fe16df68dce5a6b1d7a35dc6f702 Mon Sep 17 00:00:00 2001 From: rtemperini <29879569+rtemperini@users.noreply.github.com> Date: Mon, 7 Sep 2026 10:26:41 +0200 Subject: [PATCH 09/13] refactor(tracing): promote caller context through baggage processors Use the upstream baggage-to-span processors instead of a second request-scoped merge path, so allowlisted identity reaches every span including generate_content and survives sub-agent hops. Signed-off-by: rtemperini <29879569+rtemperini@users.noreply.github.com> Co-authored-by: Cursor --- docs/architecture/trace-context.md | 130 ++--- go/adk/pkg/a2a/executor.go | 13 +- go/adk/pkg/telemetry/attributes.go | 6 +- go/adk/pkg/telemetry/attributes_test.go | 76 +-- go/adk/pkg/telemetry/context_attributes.go | 241 +++------ .../pkg/telemetry/context_attributes_test.go | 462 ++++++------------ go/adk/pkg/telemetry/tracing.go | 13 +- go/core/internal/translator/compiler_test.go | 38 +- .../internal/translator/kagent/compiler.go | 4 +- go/core/pkg/env/otel.go | 17 +- go/go.mod | 1 + go/go.sum | 2 + .../templates/controller-deployment.yaml | 9 - .../tests/controller-deployment_test.yaml | 24 - helm/kagent/values.yaml | 27 +- .../src/kagent/adk/_agent_executor.py | 12 +- .../tests/unittests/test_agent_executor.py | 9 +- python/packages/kagent-core/pyproject.toml | 1 + .../src/kagent/core/tracing/__init__.py | 16 +- .../core/tracing/_context_attributes.py | 224 +++------ .../src/kagent/core/tracing/_utils.py | 21 +- .../tests/test_caller_context_attributes.py | 453 +++++++---------- .../tests/test_tracing_configure.py | 2 +- .../src/kagent/crewai/_executor.py | 9 +- .../kagent-crewai/tests/test_executor.py | 31 +- .../src/kagent/langgraph/_executor.py | 9 +- .../src/kagent/openai/_agent_executor.py | 9 +- .../tests/test_caller_context.py | 32 +- python/uv.lock | 26 +- 29 files changed, 616 insertions(+), 1301 deletions(-) diff --git a/docs/architecture/trace-context.md b/docs/architecture/trace-context.md index af73c4a6e0..6a3d68f2ea 100644 --- a/docs/architecture/trace-context.md +++ b/docs/architecture/trace-context.md @@ -1,7 +1,7 @@ # Caller Context in Traces Agent spans describe *what the agent did*, but they say nothing about *who asked -for it*. Kagent can promote a configurable allowlist of caller-supplied values — +for it*. Kagent can copy a configurable allowlist of caller-supplied values — an opaque user identifier, a conversation thread, a ticket ID — onto every span of a request, so traces can be filtered and grouped by the caller in Langfuse, Jaeger, Grafana Tempo, or any other OTLP backend. @@ -15,9 +15,8 @@ allowlist. | Setting | Default | Description | |---|---|---| -| Helm `otel.tracing.contextKeys` | `[]` | List of context keys or `{from, to, hash}` mappings to promote | +| Helm `otel.tracing.contextKeys` | `[]` | List of baggage keys, or `{from, to}` mappings, to copy onto spans | | Env `KAGENT_TRACE_CONTEXT_KEYS` | `""` | Comma-separated keys, or a JSON array of the same mappings | -| Helm `otel.tracing.contextHashKeySecret` | unset | Secret providing `KAGENT_TRACE_CONTEXT_HASH_KEY` for `hash: hmac-sha256` | ```yaml otel: @@ -26,7 +25,7 @@ otel: contextKeys: - {from: sub, to: user.id} - thread_id - - channel + - gen_ai.conversation.id ``` Prefer an opaque identifier such as an OIDC `sub` for `user.id`. Do not put @@ -43,26 +42,25 @@ the key and redeploy. ## Where values come from -Two sources feed the allowlist, in increasing order of precedence: +There is one path onto spans: **W3C [Baggage](https://www.w3.org/TR/baggage/)**. +Both runtimes install the upstream baggage-to-span-attribute processor +(`baggagecopy` in Go, `BaggageSpanProcessor` in Python) with a predicate that +copies only allowlisted destination keys. -| Source | Set by | Survives hops | +A2A `message.metadata` is promoted *into baggage* at the executor boundary when +the destination key is not already set. That is what lets callers that cannot +set a transport header still reach every span — including sub-agent hops — +without a second merge path or a request-scoped attribute bag. + +| Source | Set by | How it reaches spans | |---|---|---| -| W3C [Baggage](https://www.w3.org/TR/baggage/) (`baggage` header) | Any client or proxy on the request path | Yes — automatically | -| A2A `message.metadata` | The A2A caller, per message | No — one hop only | - -**Baggage is the primary mechanism.** It is the vendor-neutral OTel answer to -this problem and it needs no kagent-specific knowledge from the caller: the -controller, both runtimes, and every instrumented HTTP client already run a -composite `tracecontext + baggage` propagator, so a value set once at the edge -reaches the agent, its sub-agents, its tools, and its model calls without any -further plumbing. - -**A2A `message.metadata` is the complement** for callers that cannot set a -header — for example, a bot that speaks A2A over an SDK that exposes message -metadata but not transport headers. It is scoped to a single message, and -because it is the more specific source, it overrides baggage for the same key -when the sanitised metadata value is non-empty. A blank metadata entry cannot -wipe a baggage value. +| W3C baggage (`baggage` header) | Any client or proxy on the request path | Copied by the baggage span processor | +| A2A `message.metadata` | The A2A caller, per message | Written into baggage at the executor if the destination key is absent | + +Empty or non-scalar metadata does not wipe baggage. `{from, to}` remaps a +source key onto the destination baggage member; the processor then copies +`to`. Destination names are used as-is, so `user.id` and +`gen_ai.conversation.id` are reachable without a hardcoded unprefixed set. A key absent from both sources is simply not emitted. @@ -70,34 +68,20 @@ A key absent from both sources is simply not emitted. ## Where values land -Each mapping is read from `from` (defaulting to the entry itself) and written -as span attribute `to` (defaulting to `from`) after the prefix rules below: - ``` -baggage: sub=opaque-subject → user.id = "opaque-subject" -metadata: {"thread_id": "1717171.42"} → kagent.context.thread_id = "1717171.42" -metadata: {"channel": "C0AB1"} → kagent.context.channel = "C0AB1" +baggage: user.id=opaque-subject → user.id = "opaque-subject" +metadata: {"sub": "opaque-subject"} → user.id = "opaque-subject" # with {from: sub, to: user.id} +metadata: {"thread_id": "1717171.42"} → thread_id = "1717171.42" ``` -| Destination name | Emitted as | -|---|---| -| `user.id`, `user.hash`, `enduser.id`, `session.id` | Unprefixed (OpenTelemetry semantic conventions) | -| Anything else | `kagent.context.` | - -Invented names such as `user.asdasd` or `kagent.user_id` are prefixed. Promoted -values fill the request-scoped bag only when the key is not already set, so -caller context cannot override `kagent.user_id`, `kagent.app_name`, or -`gen_ai.*` attributes the runtime stamped. - -The attributes are merged into the **request-scoped attribute bag**, not set on -a single span. The `KagentAttributesSpanProcessor` (Python) and -`kagentAttributesSpanProcessor` (Go) stamp that bag onto every span started -during the request, so tool calls, sub-agent delegations, MCP calls, and model -calls all carry the same values. +The request-scoped kagent attribute bag still carries runtime values +(`kagent.user_id`, `gen_ai.conversation.id`, …). That processor runs after +baggagecopy, so caller baggage cannot override attributes the runtime already +stamped. -This is deliberate rather than incidental: Langfuse v4 and comparable backends -resolve trace-level filters against the attributes present on each span, so -stamping only the root span would leave most views unfilterable. +This is deliberate: Langfuse v4 and comparable backends resolve trace-level +filters against the attributes present on each span, so stamping only the +root span would leave most views unfilterable. --- @@ -107,45 +91,30 @@ stamping only the root span would leave most views unfilterable. against putting email addresses or names on telemetry at all. An OIDC `sub` is already an opaque identifier and is what `user.id` should use. -If a stable identifier must be derived from a value that itself should not -appear on a span, hash it with HMAC-SHA256 onto the registry attribute -`user.hash`: +Hashing and renaming that must happen before a value is visible on hops belong +in the authenticating gateway or the OTel Collector. Baggage travels on HTTP +headers, so a value placed there is visible to every downstream hop that +receives those headers, including model providers and HTTP MCP servers. -```yaml -contextKeys: - - {from: sub, to: user.id} - - {from: email, to: user.hash, hash: hmac-sha256} - - thread_id -``` - -`hash: hmac-sha256` requires `KAGENT_TRACE_CONTEXT_HASH_KEY` (Helm: -`otel.tracing.contextHashKeySecret`). If the key is missing, the hashed -attribute is skipped — the original value is never written onto the span. - -Hashing at promotion time only affects the span. Baggage travels on HTTP -headers, so a value placed in baggage is still visible to every downstream hop -that receives those headers, including model providers and HTTP MCP servers. -Do not put sensitive values in baggage; hash or replace them at the edge -before the request enters the cluster. +Value length is the ordinary OTel limit +(`OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT`), not a kagent-specific cap. --- ## Safety properties -Caller-supplied context is untrusted input, so promotion is constrained on every -axis: +Caller-supplied context is untrusted input, so promotion is constrained: | Risk | Control | |---|---| -| Attribute explosion / cardinality | Only allowlisted keys are read; the allowlist itself is capped at 32 entries | -| Oversized spans | Values are truncated to 256 characters, keys to 64 | -| Log or trace injection | Control characters are stripped from values | -| Shadowing semantic conventions | Custom keys are namespaced under `kagent.context.`; only `user.id`, `user.hash`, `enduser.id`, and `session.id` pass through unprefixed. Caller context fills missing keys only. | -| Leaking secrets into a trace backend | Nothing is promoted unless an operator names the key; hashed entries are omitted when the HMAC key is unset | +| Attribute explosion / cardinality | Only allowlisted keys are copied; the allowlist is capped at 32 entries | +| Oversized spans | `OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT` | +| Shadowing runtime attributes | The kagent span processor runs after baggagecopy and wins on conflict | +| Leaking secrets into a trace backend | Nothing is copied unless an operator names the key | | A tenant widening the allowlist | The allowlist is cluster-wide operator configuration; an entry of the same name in a `Harness` environment is dropped rather than inherited | Non-scalar metadata (objects, arrays) is skipped: it is unbounded in size and -meaningless as an attribute value. +meaningless as a baggage value. Choose allowlist keys deliberately. Anything named here is visible to everyone with access to the trace backend, and callers control the values. @@ -154,16 +123,16 @@ with access to the trace backend, and callers control the values. ## Renaming attributes for a backend -Some backends expect names other than the ones kagent emits. Rather than making -the attribute namespace configurable, do the rename in the OTel Collector that -already sits between kagent and the backend: +`{from, to}` is enough when the operator wants a baggage key remapped at the +executor. Further renames belong in the OTel Collector that already sits +between kagent and the backend: ```yaml processors: transform: trace_statements: - - set(span.attributes["session.id"], span.attributes["kagent.context.thread_id"]) - where span.attributes["kagent.context.thread_id"] != nil + - set(span.attributes["session.id"], span.attributes["thread_id"]) + where span.attributes["thread_id"] != nil ``` --- @@ -178,8 +147,5 @@ processors: | Python LangGraph | `python/packages/kagent-langgraph/src/kagent/langgraph/_executor.py` | | Python CrewAI | `python/packages/kagent-crewai/src/kagent/crewai/_executor.py` | | Python OpenAI | `python/packages/kagent-openai/src/kagent/openai/_agent_executor.py` | -| Controller forwarding | `go/core/v2/translator/kagent/compiler.go` | +| Controller forwarding | `go/core/internal/translator/kagent/compiler.go` | | Helm | `helm/kagent/templates/controller-configmap.yaml` | - -Both implementations share the same allowlist parsing, precedence, limits, and -sanitisation rules so the two runtimes cannot drift. diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index 7b9792f4e1..2a1b2d6cb3 100644 --- a/go/adk/pkg/a2a/executor.go +++ b/go/adk/pkg/a2a/executor.go @@ -142,16 +142,15 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.ExecutorCon if e.appName != "" { spanAttributes["kagent.app_name"] = e.appName } - // Allowlisted caller context joins the request-scoped bag rather than a - // single span, so tool, sub-agent, and model spans all carry it. - // Fill-if-absent so a caller cannot override kagent.user_id or - // gen_ai.* attributes the runtime already stamped. - telemetry.MergeCallerContextAttributes(spanAttributes, ctx, reqCtx.Message.Metadata) + // Allowlisted metadata is written into baggage so the baggagecopy + // processor stamps it on every span, including hops the runtime + // does not own. Fill-if-absent so existing baggage wins. + ctx = telemetry.ContextWithPromotedMetadata(ctx, reqCtx.Message.Metadata) ctx = telemetry.SetKAgentSpanAttributes(ctx, spanAttributes) ctx, invocationSpan := telemetry.StartInvocationSpan(ctx) defer invocationSpan.End() - // Allowlisted sources are skipped here so hashed or omitted values - // cannot leak as a2a.message.metadata. plaintext. + // Allowlisted sources are skipped here so they are not also + // stamped as a2a.message.metadata.. telemetry.SetMessageMetadataAttributes(ctx, reqCtx.Message.Metadata) e.logger.InfoContext(ctx, "execute", diff --git a/go/adk/pkg/telemetry/attributes.go b/go/adk/pkg/telemetry/attributes.go index 9426fa7258..0461e31f4b 100644 --- a/go/adk/pkg/telemetry/attributes.go +++ b/go/adk/pkg/telemetry/attributes.go @@ -91,9 +91,9 @@ func stringAttributes(attrs map[string]string) []attribute.KeyValue { // SetMessageMetadataAttributes sets scalar values from an A2A message's // metadata as span attributes on the current span. // -// Keys named in KAGENT_TRACE_CONTEXT_KEYS are skipped, including mappings -// that hash or omit the value. Stamping those sources here would put the -// original scalar on a2a.message.metadata. and bypass the allowlist. +// Keys named in KAGENT_TRACE_CONTEXT_KEYS are skipped. Those values are +// promoted into baggage and copied onto spans by baggagecopy; stamping +// them here would also emit a2a.message.metadata.. func SetMessageMetadataAttributes(ctx context.Context, metadata map[string]any) { if len(metadata) == 0 { return diff --git a/go/adk/pkg/telemetry/attributes_test.go b/go/adk/pkg/telemetry/attributes_test.go index 4b26373e42..46e790d068 100644 --- a/go/adk/pkg/telemetry/attributes_test.go +++ b/go/adk/pkg/telemetry/attributes_test.go @@ -218,79 +218,22 @@ func TestSetMessageMetadataAttributes(t *testing.T) { } } -func TestSetMessageMetadataAttributes_DoesNotStampAllowlistedHashedPlaintext(t *testing.T) { - setAllowlist(t, `[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]`) - t.Setenv(traceContextHashKeyEnvVar, "test-hmac-key") +func TestSetMessageMetadataAttributes_SkipsAllowlistedSources(t *testing.T) { + setAllowlist(t, `[{"from":"sub","to":"user.id"}]`) attrs := recordMessageMetadata(t, map[string]any{ - "email": "ada@example.com", + "sub": "opaque-subject", "channel": "C0AB1", }) - if _, exists := attrs["a2a.message.metadata.email"]; exists { - t.Fatal("allowlisted hashed source must not be stamped as a2a.message.metadata.email") + if _, exists := attrs["a2a.message.metadata.sub"]; exists { + t.Fatal("allowlisted source must not be stamped as a2a.message.metadata.sub") } - assertNoPlaintextEmail(t, attrs) if got := attrs["a2a.message.metadata.channel"].AsString(); got != "C0AB1" { t.Errorf("non-allowlisted channel = %q, want C0AB1", got) } } -func TestSetMessageMetadataAttributes_DoesNotStampWhenHMACKeyMissing(t *testing.T) { - setAllowlist(t, `[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]`) - t.Setenv(traceContextHashKeyEnvVar, "") - - attrs := recordMessageMetadata(t, map[string]any{"email": "ada@example.com"}) - - if _, exists := attrs["a2a.message.metadata.email"]; exists { - t.Fatal("omitted hashed source must not leak as a2a.message.metadata.email") - } - assertNoPlaintextEmail(t, attrs) -} - -func TestSetMessageMetadataAttributes_DoesNotStampMalformedHashSource(t *testing.T) { - setAllowlist(t, `[{"from":"email","to":"user.hash","hash":123}]`) - t.Setenv(traceContextHashKeyEnvVar, "test-hmac-key") - - attrs := recordMessageMetadata(t, map[string]any{"email": "ada@example.com"}) - - if _, exists := attrs["a2a.message.metadata.email"]; exists { - t.Fatal("malformed hash mapping must not leak plaintext as a2a.message.metadata.email") - } - assertNoPlaintextEmail(t, attrs) -} - -func TestMergeThenMessageMetadata_DoesNotLeakHashedPlaintext(t *testing.T) { - setAllowlist(t, `[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]`) - t.Setenv(traceContextHashKeyEnvVar, "test-hmac-key") - - exporter := tracetest.NewInMemoryExporter() - tp := sdktrace.NewTracerProvider( - sdktrace.WithSyncer(exporter), - sdktrace.WithSpanProcessor(kagentAttributesSpanProcessor{}), - ) - t.Cleanup(func() { _ = tp.Shutdown(context.Background()) }) - - metadata := map[string]any{"email": "ada@example.com"} - spanAttrs := map[string]string{} - MergeCallerContextAttributes(spanAttrs, context.Background(), metadata) - - tracer := tp.Tracer("test") - ctx, span := tracer.Start(context.Background(), "invocation") - ctx = SetKAgentSpanAttributes(ctx, spanAttrs) - SetMessageMetadataAttributes(ctx, metadata) - span.End() - - attrs := spanAttributesByName(t, exporter.GetSpans(), "invocation") - if _, exists := attrs["a2a.message.metadata.email"]; exists { - t.Fatal("hashed source leaked as a2a.message.metadata.email") - } - assertNoPlaintextEmail(t, attrs) - if got := attrs["user.hash"].AsString(); got == "" || strings.Contains(got, "@example.com") { - t.Errorf("user.hash = %q, want the HMAC digest", got) - } -} - func recordMessageMetadata(t *testing.T, metadata map[string]any) map[string]attribute.Value { t.Helper() exporter := tracetest.NewInMemoryExporter() @@ -313,15 +256,6 @@ func recordMessageMetadata(t *testing.T, metadata map[string]any) map[string]att return attrs } -func assertNoPlaintextEmail(t *testing.T, attrs map[string]attribute.Value) { - t.Helper() - for key, value := range attrs { - if strings.Contains(value.AsString(), "@example.com") { - t.Errorf("plaintext leaked onto %s = %q", key, value.AsString()) - } - } -} - func TestSetMessageMetadataAttributes_NilAndEmpty(t *testing.T) { exporter := tracetest.NewInMemoryExporter() tp := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter)) diff --git a/go/adk/pkg/telemetry/context_attributes.go b/go/adk/pkg/telemetry/context_attributes.go index e6c26e053a..de44622259 100644 --- a/go/adk/pkg/telemetry/context_attributes.go +++ b/go/adk/pkg/telemetry/context_attributes.go @@ -2,9 +2,6 @@ package telemetry import ( "context" - "crypto/hmac" - "crypto/sha256" - "encoding/hex" "encoding/json" "os" "strconv" @@ -13,52 +10,30 @@ import ( "unicode" "github.com/kagent-dev/kagent/go/core/pkg/env" + "go.opentelemetry.io/contrib/processors/baggagecopy" "go.opentelemetry.io/otel/baggage" ) const ( - // contextAttributePrefix namespaces custom promoted values so they cannot - // shadow a semantic convention attribute such as service.name. Only the - // explicit registry names in unprefixedAttributes pass through as-is; - // see spanAttributeName. - contextAttributePrefix = "kagent.context." - - hashHMACSHA256 = "hmac-sha256" - - maxContextKeys = 32 - maxContextKeyLength = 64 - maxContextValueLength = 256 + maxContextKeys = 32 + maxContextKeyLength = 64 ) -// Names come from env.KagentTraceContextKeys / HashKey so the controller and -// the runtime share one source of truth. -var ( - traceContextKeysEnvVar = env.KagentTraceContextKeys.Name() - traceContextHashKeyEnvVar = env.KagentTraceContextHashKey.Name() -) - -// unprefixedAttributes is the OpenTelemetry registry names callers may emit -// without a kagent.context. prefix. It is an explicit set so an invented -// name such as user.asdasd cannot masquerade as a semantic convention. -var unprefixedAttributes = map[string]struct{}{ - "user.id": {}, - "user.hash": {}, - "enduser.id": {}, - "session.id": {}, -} +// Names come from env.KagentTraceContextKeys so the controller and the +// runtime share one source of truth. +var traceContextKeysEnvVar = env.KagentTraceContextKeys.Name() // contextMapping is one allowlisted promotion: read source from baggage or -// A2A metadata and emit it as span attribute attribute (after prefix rules). +// A2A metadata and write it to baggage as attribute. The upstream +// baggagecopy processor then copies that baggage key onto every span. type contextMapping struct { source string attribute string - hash string } // contextPolicy is the parsed allowlist plus every source key named in it. -// Source keys are tracked even when a mapping is rejected (for example -// hash: 123), so generic metadata stamping cannot leak a value the operator -// asked to hash or omit. +// Source keys are skipped by SetMessageMetadataAttributes so allowlisted +// values are not also stamped as a2a.message.metadata.. type contextPolicy struct { mappings []contextMapping sources map[string]struct{} @@ -82,118 +57,74 @@ func policyCoveredContextSources() map[string]struct{} { } func parseContextPolicy() contextPolicy { - return contextPolicy{ - mappings: parseAllowedContextMappings(), - sources: parseContextPolicySources(), - } -} - -func parseContextPolicySources() map[string]struct{} { - raw := strings.TrimSpace(os.Getenv(traceContextKeysEnvVar)) - if raw == "" { - return nil - } - if strings.HasPrefix(raw, "[") { - return jsonAllowlistSourceKeys(raw) - } - sources := make(map[string]struct{}) - for key := range strings.SplitSeq(raw, ",") { - if key = strings.TrimSpace(key); key != "" { - sources[key] = struct{}{} - } - } - return sources -} - -func jsonAllowlistSourceKeys(raw string) map[string]struct{} { - var items []json.RawMessage - if err := json.Unmarshal([]byte(raw), &items); err != nil { - return nil - } - sources := make(map[string]struct{}) - for _, item := range items { - var key string - if err := json.Unmarshal(item, &key); err == nil { - if key = strings.TrimSpace(key); key != "" { - sources[key] = struct{}{} - } - continue - } - // Only bind from so a non-string hash/to cannot hide the source key. - var spec struct { - From string `json:"from"` - } - if err := json.Unmarshal(item, &spec); err != nil { - continue - } - if key := strings.TrimSpace(spec.From); key != "" { - sources[key] = struct{}{} - } + mappings := parseAllowedContextMappings() + sources := make(map[string]struct{}, len(mappings)) + for _, mapping := range mappings { + sources[mapping.source] = struct{}{} } - return sources + return contextPolicy{mappings: mappings, sources: sources} } -// CallerContextAttributes returns the caller-supplied context values that an -// operator allowlisted through KAGENT_TRACE_CONTEXT_KEYS. Merge the result into -// the request-scoped attribute bag (see SetKAgentSpanAttributes) so every span -// of the request carries them: trace-level filtering in backends such as -// Langfuse matches on each span, not only on the root. -// -// Values are read from W3C baggage first and then from the A2A message -// metadata, which is the more specific source for a single message and -// therefore wins — but only when the sanitised metadata scalar is non-empty, -// so a blank metadata entry cannot wipe a baggage value. Both are untrusted -// input, so keys must appear in the allowlist, values are stripped of control -// characters and truncated, and attribute names go through spanAttributeName. +// ContextWithPromotedMetadata copies allowlisted A2A message.metadata values +// into W3C baggage when the destination key is not already set. Existing +// baggage members named in from are remapped onto to the same way, so both +// sources share one path: the baggagecopy span processor. // -// Returns nil when the allowlist is empty, which is the default. -func CallerContextAttributes(ctx context.Context, metadata map[string]any) map[string]string { +// Empty or non-scalar metadata does not wipe baggage. Invalid baggage +// members are skipped rather than failing the request. +func ContextWithPromotedMetadata(ctx context.Context, metadata map[string]any) context.Context { mappings := allowedContextMappings() if len(mappings) == 0 { - return nil + return ctx } bag := baggage.FromContext(ctx) - attrs := make(map[string]string, len(mappings)) + changed := false for _, mapping := range mappings { - value := sanitizeContextValue(bag.Member(mapping.source).Value()) + if bag.Member(mapping.attribute).Value() != "" { + continue + } + value := "" if scalar, ok := scalarString(metadata[mapping.source]); ok { - if cleaned := sanitizeContextValue(scalar); cleaned != "" { - value = cleaned - } + value = sanitizeContextValue(scalar) + } + if value == "" && mapping.source != mapping.attribute { + value = sanitizeContextValue(bag.Member(mapping.source).Value()) } if value == "" { continue } - if mapping.hash != "" { - value = hashContextValue(value, mapping.hash) - if value == "" { - continue - } - } else { - value = truncateContextValue(value) + member, err := baggage.NewMember(mapping.attribute, value) + if err != nil { + continue } - name := spanAttributeName(mapping.attribute) - if _, exists := attrs[name]; exists { + next, err := bag.SetMember(member) + if err != nil { continue } - attrs[name] = value + bag = next + changed = true } - if len(attrs) == 0 { - return nil + if !changed { + return ctx } - return attrs + return baggage.ContextWithBaggage(ctx, bag) } -// MergeCallerContextAttributes copies promoted caller context into dst without -// replacing keys that are already set, so caller-supplied data cannot override -// kagent.user_id, kagent.app_name, or other attributes the runtime already -// stamped. -func MergeCallerContextAttributes(dst map[string]string, ctx context.Context, metadata map[string]any) { - for key, value := range CallerContextAttributes(ctx, metadata) { - if _, exists := dst[key]; !exists { - dst[key] = value - } +// AllowedBaggageCopyFilter returns a baggagecopy predicate for the +// allowlisted destination keys, or nil when promotion is off. +func AllowedBaggageCopyFilter() baggagecopy.Filter { + mappings := allowedContextMappings() + if len(mappings) == 0 { + return nil + } + allowed := make(map[string]struct{}, len(mappings)) + for _, mapping := range mappings { + allowed[mapping.attribute] = struct{}{} + } + return func(member baggage.Member) bool { + _, ok := allowed[member.Key()] + return ok } } @@ -211,7 +142,7 @@ func parseAllowedContextMappings() []contextMapping { func parseCommaAllowlist(raw string) []contextMapping { mappings := make([]contextMapping, 0, maxContextKeys) for key := range strings.SplitSeq(raw, ",") { - if mapping, ok := newContextMapping(strings.TrimSpace(key), "", ""); ok { + if mapping, ok := newContextMapping(strings.TrimSpace(key), ""); ok { mappings = append(mappings, mapping) } } @@ -227,7 +158,7 @@ func parseJSONAllowlist(raw string) []contextMapping { for _, item := range items { var key string if err := json.Unmarshal(item, &key); err == nil { - if mapping, ok := newContextMapping(key, "", ""); ok { + if mapping, ok := newContextMapping(key, ""); ok { mappings = append(mappings, mapping) } continue @@ -235,22 +166,20 @@ func parseJSONAllowlist(raw string) []contextMapping { var spec struct { From string `json:"from"` To string `json:"to"` - Hash string `json:"hash"` } if err := json.Unmarshal(item, &spec); err != nil { continue } - if mapping, ok := newContextMapping(spec.From, spec.To, spec.Hash); ok { + if mapping, ok := newContextMapping(spec.From, spec.To); ok { mappings = append(mappings, mapping) } } return mappings } -func newContextMapping(from, to, hash string) (contextMapping, bool) { +func newContextMapping(from, to string) (contextMapping, bool) { from = strings.TrimSpace(from) to = strings.TrimSpace(to) - hash = strings.TrimSpace(hash) if from == "" || len([]rune(from)) > maxContextKeyLength || !isAttributeKey(from) { return contextMapping{}, false } @@ -260,17 +189,14 @@ func newContextMapping(from, to, hash string) (contextMapping, bool) { if len([]rune(to)) > maxContextKeyLength || !isAttributeKey(to) { return contextMapping{}, false } - if hash != "" && hash != hashHMACSHA256 { - return contextMapping{}, false - } - return contextMapping{source: from, attribute: to, hash: hash}, true + return contextMapping{source: from, attribute: to}, true } func capMappings(mappings []contextMapping) []contextMapping { out := make([]contextMapping, 0, maxContextKeys) seen := make(map[string]struct{}, maxContextKeys) for _, mapping := range mappings { - id := mapping.source + "\x00" + mapping.attribute + "\x00" + mapping.hash + id := mapping.source + "\x00" + mapping.attribute if _, duplicate := seen[id]; duplicate { continue } @@ -283,20 +209,6 @@ func capMappings(mappings []contextMapping) []contextMapping { return out } -// spanAttributeName is the name written onto the span. -// -// user.id, user.hash, enduser.id, and session.id pass through unprefixed so -// operators can use the semantic convention names. Everything else is placed -// under kagent.context. so a caller-supplied service.name or kagent.user_id -// cannot shadow the real one. -func spanAttributeName(name string) string { - if _, ok := unprefixedAttributes[name]; ok { - return name - } - return contextAttributePrefix + name -} - -// isAttributeKey reports whether key is safe to use as a span attribute name. func isAttributeKey(key string) bool { return strings.IndexFunc(key, func(r rune) bool { return unicode.IsControl(r) || unicode.IsSpace(r) @@ -305,8 +217,8 @@ func isAttributeKey(key string) bool { // scalarString renders a JSON scalar from A2A message metadata as a string. // Objects and arrays are skipped: they are unbounded in size and carry no -// useful meaning as a span attribute value. JSON and protobuf Struct decode -// numbers as float64, so that is the only numeric case. +// useful meaning as a baggage value. JSON and protobuf Struct decode numbers +// as float64, so that is the only numeric case. func scalarString(value any) (string, bool) { switch v := value.(type) { case string: @@ -331,28 +243,3 @@ func sanitizeContextValue(value string) string { }, value) return strings.TrimSpace(cleaned) } - -func truncateContextValue(value string) string { - // Truncate by rune, not byte, so the limit means the same thing here as it - // does in the Python runtime. - if runes := []rune(value); len(runes) > maxContextValueLength { - return string(runes[:maxContextValueLength]) - } - return value -} - -// hashContextValue hashes value with the requested algorithm. Unknown -// algorithms and a missing HMAC key skip the attribute: never fall back to -// putting the original value on the span. -func hashContextValue(value, algorithm string) string { - if algorithm != hashHMACSHA256 { - return "" - } - key := os.Getenv(traceContextHashKeyEnvVar) - if key == "" { - return "" - } - mac := hmac.New(sha256.New, []byte(key)) - _, _ = mac.Write([]byte(value)) - return hex.EncodeToString(mac.Sum(nil)) -} diff --git a/go/adk/pkg/telemetry/context_attributes_test.go b/go/adk/pkg/telemetry/context_attributes_test.go index 3200c711f4..dcfc3a016d 100644 --- a/go/adk/pkg/telemetry/context_attributes_test.go +++ b/go/adk/pkg/telemetry/context_attributes_test.go @@ -2,13 +2,11 @@ package telemetry import ( "context" - "crypto/hmac" - "crypto/sha256" - "encoding/hex" "strconv" "strings" "testing" + "go.opentelemetry.io/contrib/processors/baggagecopy" "go.opentelemetry.io/otel/baggage" sdktrace "go.opentelemetry.io/otel/sdk/trace" "go.opentelemetry.io/otel/sdk/trace/tracetest" @@ -39,427 +37,239 @@ func baggageContext(t *testing.T, members map[string]string) context.Context { return baggage.ContextWithBaggage(context.Background(), bag) } -func hmacSHA256Hex(t *testing.T, key, value string) string { - t.Helper() - mac := hmac.New(sha256.New, []byte(key)) - _, _ = mac.Write([]byte(value)) - return hex.EncodeToString(mac.Sum(nil)) +func baggageValues(ctx context.Context) map[string]string { + out := map[string]string{} + for _, member := range baggage.FromContext(ctx).Members() { + out[member.Key()] = member.Value() + } + return out } -func TestCallerContextAttributes(t *testing.T) { +func TestContextWithPromotedMetadata(t *testing.T) { tests := []struct { name string allowlist string - hashKey string baggageVals map[string]string metadata map[string]any want map[string]string }{ { - name: "empty allowlist disables promotion", - allowlist: "", - baggageVals: map[string]string{"sub": "opaque-subject"}, - metadata: map[string]any{"thread_id": "T123"}, - want: nil, - }, - { - name: "promotes allowlisted baggage", - allowlist: "sub,thread_id", - baggageVals: map[string]string{"sub": "opaque-subject", "thread_id": "T123"}, - want: map[string]string{ - "kagent.context.sub": "opaque-subject", - "kagent.context.thread_id": "T123", - }, + name: "empty allowlist is a no-op", + baggageVals: map[string]string{"user.id": "opaque-subject"}, + metadata: map[string]any{"user.id": "from-metadata"}, + want: map[string]string{"user.id": "opaque-subject"}, }, { - name: "promotes allowlisted message metadata", - allowlist: "thread_id,channel", - metadata: map[string]any{"thread_id": "1717171.4242", "channel": "C0AB1"}, - want: map[string]string{ - "kagent.context.thread_id": "1717171.4242", - "kagent.context.channel": "C0AB1", - }, + name: "leaves existing baggage in place", + allowlist: "user.id", + baggageVals: map[string]string{"user.id": "from-baggage"}, + metadata: map[string]any{"user.id": "from-metadata"}, + want: map[string]string{"user.id": "from-baggage"}, }, { - name: "message metadata overrides baggage", - allowlist: "sub", - baggageVals: map[string]string{"sub": "from-baggage"}, - metadata: map[string]any{"sub": "from-metadata"}, - want: map[string]string{"kagent.context.sub": "from-metadata"}, + name: "promotes allowlisted metadata into baggage", + allowlist: "user.id,thread_id", + metadata: map[string]any{"user.id": "opaque-subject", "thread_id": "T123", "secret": "nope"}, + want: map[string]string{"user.id": "opaque-subject", "thread_id": "T123"}, }, { - name: "empty metadata does not override baggage", - allowlist: "sub", - baggageVals: map[string]string{"sub": "from-baggage"}, - metadata: map[string]any{"sub": " "}, - want: map[string]string{"kagent.context.sub": "from-baggage"}, + name: "empty metadata does not wipe baggage", + allowlist: "user.id", + baggageVals: map[string]string{"user.id": "from-baggage"}, + metadata: map[string]any{"user.id": " "}, + want: map[string]string{"user.id": "from-baggage"}, }, { - name: "ignores keys outside the allowlist", - allowlist: "thread_id", - baggageVals: map[string]string{"secret.token": "s3cret"}, - metadata: map[string]any{"thread_id": "T1", "extra": "nope"}, - want: map[string]string{"kagent.context.thread_id": "T1"}, + name: "remaps baggage from onto to when to is absent", + allowlist: `[{"from":"sub","to":"user.id"}]`, + baggageVals: map[string]string{"sub": "opaque-subject"}, + want: map[string]string{"sub": "opaque-subject", "user.id": "opaque-subject"}, }, { - name: "renders scalar metadata types", - allowlist: "count,ratio,enabled", - metadata: map[string]any{ - "count": float64(7), - "ratio": float64(2.5), - "enabled": true, - }, - want: map[string]string{ - "kagent.context.count": "7", - "kagent.context.ratio": "2.5", - "kagent.context.enabled": "true", - }, + name: "does not remap over an existing to", + allowlist: `[{"from":"sub","to":"user.id"}]`, + baggageVals: map[string]string{"sub": "from-sub", "user.id": "already"}, + want: map[string]string{"sub": "from-sub", "user.id": "already"}, }, { - name: "formats large floats without scientific notation", - allowlist: "value", - metadata: map[string]any{"value": float64(1e20)}, - want: map[string]string{"kagent.context.value": "100000000000000000000"}, + name: "maps metadata from onto to", + allowlist: `[{"from":"sub","to":"user.id"},"gen_ai.conversation.id"]`, + metadata: map[string]any{"sub": "opaque-subject", "gen_ai.conversation.id": "sess-1"}, + want: map[string]string{"user.id": "opaque-subject", "gen_ai.conversation.id": "sess-1"}, }, { - name: "formats tiny floats without rounding to zero", + name: "skips non-scalars", allowlist: "value", - metadata: map[string]any{"value": float64(1e-20)}, - want: map[string]string{"kagent.context.value": "0.00000000000000000001"}, + metadata: map[string]any{"value": map[string]any{"nested": true}}, + want: map[string]string{}, }, { - name: "skips non-scalar and empty metadata values", - allowlist: "nested,list,blank", - metadata: map[string]any{ - "nested": map[string]any{"a": "b"}, - "list": []string{"a"}, - "blank": "", - }, - want: nil, + name: "renders scalar metadata", + allowlist: "flag,count", + metadata: map[string]any{"flag": true, "count": float64(3)}, + want: map[string]string{"flag": "true", "count": "3"}, }, { name: "strips control characters", allowlist: "note", metadata: map[string]any{"note": "line\nbreak\tand\x00nul"}, - want: map[string]string{"kagent.context.note": "linebreakandnul"}, - }, - { - name: "ignores allowlist entries that are not valid attribute keys", - allowlist: "good, bad key ,\tanother\tbad", - metadata: map[string]any{"good": "yes", "bad key": "no"}, - want: map[string]string{"kagent.context.good": "yes"}, - }, - { - // Only the explicit registry set passes through unprefixed. - name: "registry attributes stay unprefixed", - allowlist: "user.id,enduser.id,session.id,channel", - metadata: map[string]any{ - "user.id": "opaque-subject", - "enduser.id": "end-user", - "session.id": "sess-1", - "channel": "C0AB1", - }, - want: map[string]string{ - "user.id": "opaque-subject", - "enduser.id": "end-user", - "session.id": "sess-1", - "kagent.context.channel": "C0AB1", - }, - }, - { - name: "unknown user attributes are prefixed", - allowlist: `[{"from":"x","to":"user.asdasd"},{"from":"y","to":"user.email"}]`, - metadata: map[string]any{"x": "nope", "y": "also-nope"}, - want: map[string]string{ - "kagent.context.user.asdasd": "nope", - "kagent.context.user.email": "also-nope", - }, - }, - { - // session.id is the registry name; other session.* keys are not. - name: "session.id is unprefixed but session.foo is not", - allowlist: "session.id,session.foo", - metadata: map[string]any{"session.id": "sess-1", "session.foo": "other"}, - want: map[string]string{ - "session.id": "sess-1", - "kagent.context.session.foo": "other", - }, - }, - { - name: "maps source keys onto registry names", - allowlist: `[{"from":"sub","to":"user.id"},{"from":"thread_id","to":"session.id"},"channel"]`, - metadata: map[string]any{ - "sub": "opaque-subject", - "thread_id": "T123", - "channel": "C0AB1", - }, - want: map[string]string{ - "user.id": "opaque-subject", - "session.id": "T123", - "kagent.context.channel": "C0AB1", - }, - }, - { - name: "kagent destination names are prefixed", - allowlist: `[{"from":"uid","to":"kagent.user_id"},{"from":"tid","to":"kagent.thread_id"}]`, - metadata: map[string]any{"uid": "attacker", "tid": "T123"}, - want: map[string]string{ - "kagent.context.kagent.user_id": "attacker", - "kagent.context.kagent.thread_id": "T123", - }, - }, - { - name: "invalid JSON allowlist promotes nothing", - allowlist: `[{"from":"sub"`, - metadata: map[string]any{"sub": "opaque-subject"}, - want: nil, - }, - { - name: "non-string hash does not emit plaintext", - allowlist: `[{"from":"email","to":"user.hash","hash":123}]`, - hashKey: "test-hmac-key", - metadata: map[string]any{"email": "ada@example.com"}, - want: nil, + want: map[string]string{"note": "linebreakandnul"}, }, { - name: "non-string to drops the mapping", - allowlist: `[{"from":"thread_id","to":123}]`, - metadata: map[string]any{"thread_id": "T1"}, - want: nil, + name: "drops keys with whitespace", + allowlist: "bad key,user.id", + metadata: map[string]any{"bad key": "x", "user.id": "ok"}, + want: map[string]string{"user.id": "ok"}, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { setAllowlist(t, tt.allowlist) - if tt.hashKey != "" { - t.Setenv(traceContextHashKeyEnvVar, tt.hashKey) - } - - got := CallerContextAttributes(baggageContext(t, tt.baggageVals), tt.metadata) - + got := baggageValues(ContextWithPromotedMetadata(baggageContext(t, tt.baggageVals), tt.metadata)) if len(got) != len(tt.want) { - t.Fatalf("got %v, want %v", got, tt.want) + t.Fatalf("got %#v, want %#v", got, tt.want) } for key, want := range tt.want { if got[key] != want { - t.Errorf("%s = %q, want %q", key, got[key], want) + t.Errorf("%s: got %q, want %q", key, got[key], want) } } }) } } -func TestMergeCallerContextAttributes_DoesNotOverrideExisting(t *testing.T) { - setAllowlist(t, `[{"from":"sub","to":"user.id"},"thread_id"]`) - - dst := map[string]string{ - "user.id": "runtime-user", - "kagent.user_id": "runtime-user", - } - MergeCallerContextAttributes(dst, context.Background(), map[string]any{ - "sub": "attacker", - "thread_id": "T1", - }) - - if dst["user.id"] != "runtime-user" { - t.Errorf("user.id = %q, want runtime-user", dst["user.id"]) - } - if dst["kagent.user_id"] != "runtime-user" { - t.Errorf("kagent.user_id = %q, want runtime-user", dst["kagent.user_id"]) - } - if dst["kagent.context.thread_id"] != "T1" { - t.Errorf("thread_id = %q, want T1", dst["kagent.context.thread_id"]) - } -} - -func TestCallerContextAttributes_HashesWithHMACSHA256(t *testing.T) { - const key = "test-hmac-key" - setAllowlist(t, `[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]`) - t.Setenv(traceContextHashKeyEnvVar, key) - - got := CallerContextAttributes(context.Background(), map[string]any{ - "email": "ada@example.com", - }) - - want := hmacSHA256Hex(t, key, "ada@example.com") - if got["user.hash"] != want { - t.Errorf("user.hash = %q, want %q", got["user.hash"], want) +func TestContextWithPromotedMetadata_CapsAllowlist(t *testing.T) { + keys := make([]string, maxContextKeys+4) + metadata := map[string]any{} + for i := range keys { + keys[i] = "k" + strconv.Itoa(i) + metadata[keys[i]] = "v" } - for name, value := range got { - if strings.Contains(value, "@example.com") { - t.Errorf("plaintext leaked onto %s", name) - } - } -} - -func TestCallerContextAttributes_HashWithoutKeyEmitsNothing(t *testing.T) { - // Missing HMAC key must not fall back to putting the original value on the span. - setAllowlist(t, `[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]`) - t.Setenv(traceContextHashKeyEnvVar, "") - - got := CallerContextAttributes(context.Background(), map[string]any{ - "email": "ada@example.com", - }) - if got != nil { - t.Errorf("got %v, want nothing when the HMAC key is unset", got) + setAllowlist(t, strings.Join(keys, ",")) + got := baggageValues(ContextWithPromotedMetadata(context.Background(), metadata)) + if len(got) != maxContextKeys { + t.Fatalf("promoted %d keys, want %d", len(got), maxContextKeys) } } -func TestCallerContextAttributes_UnknownHashEmitsNothing(t *testing.T) { - setAllowlist(t, `[{"from":"email","to":"user.hash","hash":"md5"}]`) - t.Setenv(traceContextHashKeyEnvVar, "test-hmac-key") - - got := CallerContextAttributes(context.Background(), map[string]any{ - "email": "ada@example.com", - }) - if got != nil { - t.Errorf("got %v, want nothing for an unsupported hash", got) +func TestAllowedBaggageCopyFilter(t *testing.T) { + if AllowedBaggageCopyFilter() != nil { + t.Fatal("empty allowlist must not install a baggagecopy filter") } -} -func TestCallerContextAttributes_TruncatesLongValues(t *testing.T) { - setAllowlist(t, "note") - - got := CallerContextAttributes(context.Background(), map[string]any{ - "note": strings.Repeat("a", maxContextValueLength*2), - }) - - if len(got["kagent.context.note"]) != maxContextValueLength { - t.Errorf("value length = %d, want %d", len(got["kagent.context.note"]), maxContextValueLength) + setAllowlist(t, `[{"from":"sub","to":"user.id"},"thread_id"]`) + filter := AllowedBaggageCopyFilter() + if filter == nil { + t.Fatal("expected a filter") } -} - -func TestAllowedContextMappings_CapsListLength(t *testing.T) { - keys := make([]string, 0, maxContextKeys*2) - for i := range maxContextKeys * 2 { - keys = append(keys, "key"+strconv.Itoa(i)) + if !filter(mustMember(t, "user.id", "x")) { + t.Error("user.id should be copied") } - setAllowlist(t, strings.Join(keys, ",")) - - if got := len(allowedContextMappings()); got != maxContextKeys { - t.Errorf("allowlist length = %d, want %d", got, maxContextKeys) + if !filter(mustMember(t, "thread_id", "x")) { + t.Error("thread_id should be copied") } -} - -func TestAllowedContextMappings_DropsOverLongAndDuplicateKeys(t *testing.T) { - setAllowlist(t, "a,a,"+strings.Repeat("b", maxContextKeyLength+1)+",c") - - got := allowedContextMappings() - - want := []string{"a", "c"} - if len(got) != len(want) { - t.Fatalf("got %#v, want %v", got, want) + if filter(mustMember(t, "sub", "x")) { + t.Error("source-only key sub should not be copied") } - for i, key := range want { - if got[i].source != key { - t.Errorf("key %d = %q, want %q", i, got[i].source, key) - } + if filter(mustMember(t, "secret", "x")) { + t.Error("non-allowlisted key should not be copied") } } -func TestAllowedContextMappings_CountsKeyLengthInRunes(t *testing.T) { - kept := strings.Repeat("键", 40) - dropped := strings.Repeat("键", 65) - setAllowlist(t, kept+","+dropped) - - got := allowedContextMappings() - if len(got) != 1 || got[0].source != kept { - t.Fatalf("got %#v, want the 40-rune CJK key kept and the 65-rune key dropped", got) - } -} +func TestCallerContextLandsOnEverySpan(t *testing.T) { + setAllowlist(t, "user.id") -func TestAllowedContextMappings_CachesUntilReset(t *testing.T) { - setAllowlist(t, "a") - if got := allowedContextMappings(); len(got) != 1 || got[0].source != "a" { - t.Fatalf("got %#v, want [a]", got) - } + exporter := tracetest.NewInMemoryExporter() + tp := sdktrace.NewTracerProvider( + sdktrace.WithSyncer(exporter), + sdktrace.WithSpanProcessor(baggagecopy.NewSpanProcessor(AllowedBaggageCopyFilter())), + sdktrace.WithSpanProcessor(kagentAttributesSpanProcessor{}), + ) + t.Cleanup(func() { _ = tp.Shutdown(context.Background()) }) - t.Setenv(traceContextKeysEnvVar, "b") - if got := allowedContextMappings(); len(got) != 1 || got[0].source != "a" { - t.Errorf("OnceValue lost the first parse: %#v", got) - } + ctx := ContextWithPromotedMetadata(baggageContext(t, map[string]string{"user.id": "opaque-subject"}), nil) + ctx = SetKAgentSpanAttributes(ctx, map[string]string{"kagent.user_id": "runtime-user"}) + tracer := tp.Tracer("test") + ctx, root := tracer.Start(ctx, "invocation") + _, tool := tracer.Start(ctx, "tool") + _, model := tracer.Start(ctx, "generate_content") + model.End() + tool.End() + root.End() - resetAllowedContextMappings() - if got := allowedContextMappings(); len(got) != 1 || got[0].source != "b" { - t.Errorf("reset did not reparse: %#v", got) + for _, name := range []string{"invocation", "tool", "generate_content"} { + attrs := spanAttributesByName(t, exporter.GetSpans(), name) + if got := attrs["user.id"].AsString(); got != "opaque-subject" { + t.Errorf("%s user.id = %q, want opaque-subject", name, got) + } + if got := attrs["kagent.user_id"].AsString(); got != "runtime-user" { + t.Errorf("%s kagent.user_id = %q, want runtime-user", name, got) + } } } -// Langfuse and comparable backends filter on attributes present on each span, -// so the promoted values must reach descendants, not just the root span. -func TestCallerContextAttributes_ReachEverySpan(t *testing.T) { +func TestMetadataPromotionReachesGenerateContent(t *testing.T) { setAllowlist(t, `[{"from":"sub","to":"user.id"}]`) exporter := tracetest.NewInMemoryExporter() tp := sdktrace.NewTracerProvider( sdktrace.WithSyncer(exporter), + sdktrace.WithSpanProcessor(baggagecopy.NewSpanProcessor(AllowedBaggageCopyFilter())), sdktrace.WithSpanProcessor(kagentAttributesSpanProcessor{}), ) - t.Cleanup(func() { - _ = tp.Shutdown(context.Background()) - }) - - ctx := baggageContext(t, map[string]string{"sub": "opaque-subject"}) - ctx = SetKAgentSpanAttributes(ctx, CallerContextAttributes(ctx, nil)) + t.Cleanup(func() { _ = tp.Shutdown(context.Background()) }) + ctx := ContextWithPromotedMetadata(context.Background(), map[string]any{"sub": "opaque-subject"}) + ctx = SetKAgentSpanAttributes(ctx, map[string]string{"gen_ai.conversation.id": "sess-runtime"}) tracer := tp.Tracer("test") - ctx, root := tracer.Start(ctx, "root") - ctx, tool := tracer.Start(ctx, "execute_tool") + ctx, root := tracer.Start(ctx, "invocation") _, model := tracer.Start(ctx, "generate_content") model.End() - tool.End() root.End() - spans := exporter.GetSpans() - if len(spans) != 3 { - t.Fatalf("expected 3 spans, got %d", len(spans)) - } - for _, name := range []string{"root", "execute_tool", "generate_content"} { - attrs := spanAttributesByName(t, spans, name) + for _, name := range []string{"invocation", "generate_content"} { + attrs := spanAttributesByName(t, exporter.GetSpans(), name) if got := attrs["user.id"].AsString(); got != "opaque-subject" { - t.Errorf("span %q: user.id = %q, want %q", name, got, "opaque-subject") + t.Errorf("%s user.id = %q, want opaque-subject", name, got) + } + if got := attrs["gen_ai.conversation.id"].AsString(); got != "sess-runtime" { + t.Errorf("%s gen_ai.conversation.id = %q, want sess-runtime", name, got) + } + if _, exists := attrs["a2a.message.metadata.sub"]; exists { + t.Errorf("%s unexpectedly stamped a2a.message.metadata.sub", name) } } } -func TestCallerContextAttributes_DisabledLeavesSpansUnchanged(t *testing.T) { - setAllowlist(t, "") +func TestRuntimeAttributesWinOverAllowlistedBaggage(t *testing.T) { + setAllowlist(t, "gen_ai.conversation.id") exporter := tracetest.NewInMemoryExporter() tp := sdktrace.NewTracerProvider( sdktrace.WithSyncer(exporter), + sdktrace.WithSpanProcessor(baggagecopy.NewSpanProcessor(AllowedBaggageCopyFilter())), sdktrace.WithSpanProcessor(kagentAttributesSpanProcessor{}), ) - t.Cleanup(func() { - _ = tp.Shutdown(context.Background()) - }) + t.Cleanup(func() { _ = tp.Shutdown(context.Background()) }) - ctx := baggageContext(t, map[string]string{"sub": "opaque-subject"}) - ctx = SetKAgentSpanAttributes(ctx, CallerContextAttributes(ctx, map[string]any{"thread_id": "T1"})) - - _, span := tp.Tracer("test").Start(ctx, "root") + ctx := baggageContext(t, map[string]string{"gen_ai.conversation.id": "from-caller"}) + ctx = SetKAgentSpanAttributes(ctx, map[string]string{"gen_ai.conversation.id": "from-runtime"}) + tracer := tp.Tracer("test") + _, span := tracer.Start(ctx, "generate_content") span.End() - for _, attr := range exporter.GetSpans()[0].Attributes { - key := string(attr.Key) - if strings.HasPrefix(key, contextAttributePrefix) || key == "user.id" { - t.Errorf("unexpected promoted attribute %q", attr.Key) - } + attrs := spanAttributesByName(t, exporter.GetSpans(), "generate_content") + if got := attrs["gen_ai.conversation.id"].AsString(); got != "from-runtime" { + t.Fatalf("gen_ai.conversation.id = %q, want runtime value to win", got) } } -// Custom keys still cannot shadow a semantic convention attribute such as -// service.name. Registry names are the documented exception. -func TestCallerContextAttributes_CannotShadowSemanticConventions(t *testing.T) { - setAllowlist(t, "service.name") - - got := CallerContextAttributes(context.Background(), map[string]any{"service.name": "impostor"}) - - if _, shadowed := got["service.name"]; shadowed { - t.Error("service.name must not be settable by a caller") - } - if got["kagent.context.service.name"] != "impostor" { - t.Errorf("got %v, want the value namespaced under %q", got, contextAttributePrefix) +func mustMember(t *testing.T, key, value string) baggage.Member { + t.Helper() + member, err := baggage.NewMember(key, value) + if err != nil { + t.Fatal(err) } + return member } diff --git a/go/adk/pkg/telemetry/tracing.go b/go/adk/pkg/telemetry/tracing.go index 3a1d22a936..3ca7eed224 100644 --- a/go/adk/pkg/telemetry/tracing.go +++ b/go/adk/pkg/telemetry/tracing.go @@ -9,6 +9,7 @@ import ( "strings" "time" + "go.opentelemetry.io/contrib/processors/baggagecopy" "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploggrpc" "go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploghttp" @@ -195,11 +196,17 @@ func newTracerProvider(ctx context.Context, res *resource.Resource) (*sdktrace.T return nil, err } - return sdktrace.NewTracerProvider( - sdktrace.WithSpanProcessor(kagentAttributesSpanProcessor{}), + opts := []sdktrace.TracerProviderOption{ sdktrace.WithBatcher(exporter), sdktrace.WithResource(res), - ), nil + } + // Baggagecopy runs first so the request-scoped bag can still win on + // keys the runtime already stamped (kagent.user_id, gen_ai.*). + if filter := AllowedBaggageCopyFilter(); filter != nil { + opts = append(opts, sdktrace.WithSpanProcessor(baggagecopy.NewSpanProcessor(filter))) + } + opts = append(opts, sdktrace.WithSpanProcessor(kagentAttributesSpanProcessor{})) + return sdktrace.NewTracerProvider(opts...), nil } func newLoggerProvider(ctx context.Context, res *resource.Resource) (*sdklog.LoggerProvider, error) { diff --git a/go/core/internal/translator/compiler_test.go b/go/core/internal/translator/compiler_test.go index 97bf2d0884..f7bc1392d9 100644 --- a/go/core/internal/translator/compiler_test.go +++ b/go/core/internal/translator/compiler_test.go @@ -375,10 +375,9 @@ func TestCompileAgentTemplateSharedAgent(t *testing.T) { require.Contains(t, string(revision.Provenance), `"name":"researcher"`) } -// KAGENT_TRACE_CONTEXT_KEYS and KAGENT_TRACE_CONTEXT_HASH_KEY are not OTEL_ -// variables, so OtelEnvFromProcess does not carry them and they need -// forwarding of their own. They are also operator policy, so a Harness must -// not be able to widen, enable, or supply the HMAC key. +// KAGENT_TRACE_CONTEXT_KEYS is not an OTEL_ variable, so OtelEnvFromProcess +// does not carry it and it needs forwarding of its own. It is also operator +// policy, so a Harness must not be able to widen or enable it. func TestCompileAgentTemplateForwardsTraceContextPolicy(t *testing.T) { template := &v1alpha3.AgentTemplate{ ObjectMeta: metav1.ObjectMeta{Name: "helper", Namespace: "test"}, @@ -389,46 +388,31 @@ func TestCompileAgentTemplateForwardsTraceContextPolicy(t *testing.T) { tests := []struct { name string keys string - hashKey string harnessEnv []v1alpha3.HarnessEnvVar wantKeys string - wantHash string }{ {name: "absent when unconfigured"}, - {name: "keys forwarded when configured", keys: "sub,thread_id", wantKeys: "sub,thread_id"}, + {name: "keys forwarded when configured", keys: "user.id,thread_id", wantKeys: "user.id,thread_id"}, { - name: "hash key forwarded when configured", + name: "mappings forwarded when configured", keys: `[{"from":"sub","to":"user.id"}]`, - hashKey: "test-hmac-key", wantKeys: `[{"from":"sub","to":"user.id"}]`, - wantHash: "test-hmac-key", }, { name: "harness cannot widen the allowlist", - keys: "sub", + keys: "user.id", harnessEnv: []v1alpha3.HarnessEnvVar{{Name: env.KagentTraceContextKeys.Name(), Value: &harnessSupplied}}, - wantKeys: "sub", + wantKeys: "user.id", }, { name: "harness cannot enable promotion", harnessEnv: []v1alpha3.HarnessEnvVar{{Name: env.KagentTraceContextKeys.Name(), Value: &harnessSupplied}}, }, - { - name: "harness cannot supply the HMAC key", - harnessEnv: []v1alpha3.HarnessEnvVar{{Name: env.KagentTraceContextHashKey.Name(), Value: &harnessSupplied}}, - }, - { - name: "harness cannot replace the HMAC key", - hashKey: "operator-hmac-key", - harnessEnv: []v1alpha3.HarnessEnvVar{{Name: env.KagentTraceContextHashKey.Name(), Value: &harnessSupplied}}, - wantHash: "operator-hmac-key", - }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Setenv(env.KagentTraceContextKeys.Name(), tt.keys) - t.Setenv(env.KagentTraceContextHashKey.Name(), tt.hashKey) harness := &v1alpha3.Harness{ ObjectMeta: metav1.ObjectMeta{Name: "kagent", Namespace: "test"}, Spec: v1alpha3.HarnessSpec{ @@ -446,19 +430,13 @@ func TestCompileAgentTemplateForwardsTraceContextPolicy(t *testing.T) { require.NoError(t, err) gotKeys, seenKeys := "", 0 - gotHash, seenHash := "", 0 for _, variable := range revision.Environment { - switch variable.Name { - case env.KagentTraceContextKeys.Name(): + if variable.Name == env.KagentTraceContextKeys.Name() { gotKeys, seenKeys = variable.Value, seenKeys+1 - case env.KagentTraceContextHashKey.Name(): - gotHash, seenHash = variable.Value, seenHash+1 } } require.LessOrEqual(t, seenKeys, 1, "environment must not contain a duplicate keys entry") - require.LessOrEqual(t, seenHash, 1, "environment must not contain a duplicate hash key entry") require.Equal(t, tt.wantKeys, gotKeys) - require.Equal(t, tt.wantHash, gotHash) }) } } diff --git a/go/core/internal/translator/kagent/compiler.go b/go/core/internal/translator/kagent/compiler.go index 23c6571371..f528db2b8a 100644 --- a/go/core/internal/translator/kagent/compiler.go +++ b/go/core/internal/translator/kagent/compiler.go @@ -76,10 +76,8 @@ func (c *Compiler) Compile(ctx context.Context, input *v2translator.HarnessInput // Which caller-supplied context reaches traces is cluster-wide operator // policy, so a Harness must be able to neither widen nor enable it. Dropping // any inherited entry before applying the operator's value is what makes that - // hold when the operator has configured nothing at all. The HMAC key is the - // same class of policy: a tenant must not supply it. + // hold when the operator has configured nothing at all. environment = applyOperatorOnlyEnv(environment, env.KagentTraceContextKeys) - environment = applyOperatorOnlyEnv(environment, env.KagentTraceContextHashKey) environment = append(environment, v2translator.OtelEnvFromProcess()...) environment = adkconfig.DedupeEnv(environment) provenance, err := c.config.BuildProvenance(ctx, harness, compiled.Templates, compiled.Models, environment) diff --git a/go/core/pkg/env/otel.go b/go/core/pkg/env/otel.go index 4145a729be..2e77fe05a1 100644 --- a/go/core/pkg/env/otel.go +++ b/go/core/pkg/env/otel.go @@ -41,19 +41,10 @@ var ( KagentTraceContextKeys = RegisterStringVar( "KAGENT_TRACE_CONTEXT_KEYS", "", - "Allowlist of caller-supplied context keys promoted onto every agent span. "+ - "Accepts a comma-separated list of source keys, or a JSON array of strings and "+ - "{from, to, hash} objects. Registry names (user.id, user.hash, enduser.id, session.id) "+ - "are left unprefixed; everything else is emitted as kagent.context.. "+ - "Empty (the default) disables promotion.", - ComponentAgentRuntime, - ) - - KagentTraceContextHashKey = RegisterStringVar( - "KAGENT_TRACE_CONTEXT_HASH_KEY", - "", - "HMAC-SHA256 key used when a KAGENT_TRACE_CONTEXT_KEYS mapping sets hash: hmac-sha256. "+ - "Hashed attributes are omitted when this is unset, rather than emitting the original value.", + "Allowlist of caller-supplied context keys copied from W3C baggage onto every agent span. "+ + "Accepts a comma-separated list of keys, or a JSON array of strings and {from, to} "+ + "objects. Destination names are used as-is; renaming and hashing belong in the "+ + "collector or gateway. Empty (the default) disables promotion.", ComponentAgentRuntime, ) ) diff --git a/go/go.mod b/go/go.mod index 3bef8daf6d..5acf857b50 100644 --- a/go/go.mod +++ b/go/go.mod @@ -52,6 +52,7 @@ require ( go.opentelemetry.io/contrib/exporters/autoexport v0.69.0 go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.69.0 go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.69.0 + go.opentelemetry.io/contrib/processors/baggagecopy v0.16.1 go.opentelemetry.io/otel v1.44.0 go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploggrpc v0.20.0 go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploghttp v0.20.0 diff --git a/go/go.sum b/go/go.sum index 7770285608..2336b86481 100644 --- a/go/go.sum +++ b/go/go.sum @@ -1222,6 +1222,8 @@ go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.6 go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.69.0/go.mod h1:D7J12YRapIekYyPWgGPlA/23pRmpSEZC5xJC/TTLI9U= go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.69.0 h1:8tvICD4vSTOOsNrsI4Ljf6C+6UKvpTEH5XY3JMoyPoo= go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.69.0/go.mod h1:z9+yiacE0IHRqM4qFfkbt/JYlmYXgss8GY/jXoNuPJI= +go.opentelemetry.io/contrib/processors/baggagecopy v0.16.1 h1:yIjm4hrmlS2tXrXvjvQ0FkMC5gM3MHbPCvStadgF1nc= +go.opentelemetry.io/contrib/processors/baggagecopy v0.16.1/go.mod h1:5PIbXha494ZMsloI0o5HI9/7qv/YPcZ9eNrwr94vUeI= go.opentelemetry.io/otel v1.44.0 h1:JjwHmHpA4iZ3wBxluu2fbbE7j4kqlE8jXyAyPXH7HqU= go.opentelemetry.io/otel v1.44.0/go.mod h1:BMgjTHL9WPRlRjL2oZCBTL4whCGtXch2H4BhOPIAyYc= go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploggrpc v0.20.0 h1:rydZ9sxbcFdm/oWrVyfLTjHIygMgv0bEeMd+3B/BvoM= diff --git a/helm/kagent/templates/controller-deployment.yaml b/helm/kagent/templates/controller-deployment.yaml index 4a6e960def..8859e4725f 100644 --- a/helm/kagent/templates/controller-deployment.yaml +++ b/helm/kagent/templates/controller-deployment.yaml @@ -139,15 +139,6 @@ spec: {{- with .Values.controller.env }} {{- toYaml . | nindent 12 }} {{- end }} - {{- with .Values.otel.tracing.contextHashKeySecret }} - {{- if .name }} - - name: KAGENT_TRACE_CONTEXT_HASH_KEY - valueFrom: - secretKeyRef: - name: {{ .name | quote }} - key: {{ .key | default "hmac-key" | quote }} - {{- end }} - {{- end }} {{- if and .Values.controller.substrate .Values.controller.substrate.enabled }} - name: SUBSTRATE_ATE_API_ENDPOINT value: {{ include "kagent.substrate.ateApiEndpoint" . | quote }} diff --git a/helm/kagent/tests/controller-deployment_test.yaml b/helm/kagent/tests/controller-deployment_test.yaml index 576da4fd66..ba2b7e75f9 100644 --- a/helm/kagent/tests/controller-deployment_test.yaml +++ b/helm/kagent/tests/controller-deployment_test.yaml @@ -835,27 +835,3 @@ tests: - equal: path: data.KAGENT_TRACE_CONTEXT_KEYS value: '[{"from":"sub","to":"user.id"},{"from":"thread_id","to":"session.id"},"channel"]' - - - it: should omit KAGENT_TRACE_CONTEXT_HASH_KEY by default - template: controller-deployment.yaml - asserts: - - notContains: - path: spec.template.spec.containers[0].env - content: - name: KAGENT_TRACE_CONTEXT_HASH_KEY - - - it: should inject KAGENT_TRACE_CONTEXT_HASH_KEY from the configured secret - template: controller-deployment.yaml - set: - otel.tracing.contextHashKeySecret: - name: trace-context-hmac - key: hmac-key - asserts: - - contains: - path: spec.template.spec.containers[0].env - content: - name: KAGENT_TRACE_CONTEXT_HASH_KEY - valueFrom: - secretKeyRef: - name: trace-context-hmac - key: hmac-key diff --git a/helm/kagent/values.yaml b/helm/kagent/values.yaml index 311a5cdc27..be766d652e 100644 --- a/helm/kagent/values.yaml +++ b/helm/kagent/values.yaml @@ -838,30 +838,21 @@ oauth2-proxy: otel: tracing: enabled: false - # Allowlist of caller-supplied context keys promoted onto every agent span. - # Values are read from W3C baggage and A2A message metadata. Empty (the - # default) disables promotion. + # Allowlist of caller-supplied context keys copied from W3C baggage onto + # every agent span. A2A message.metadata values for the same keys are + # written into baggage at the executor so they survive sub-agent hops. + # Empty (the default) disables promotion. # - # Entries may be a source key or a mapping. Prefer an opaque identifier - # such as an OIDC subject for user.id; do not put names or email addresses - # on spans. Registry names (user.id, user.hash, enduser.id, session.id) - # are left unprefixed; everything else is emitted as kagent.context.. + # Entries may be a key or a {from, to} mapping. Destination names are + # used as-is. Prefer an opaque identifier such as an OIDC subject for + # user.id; do not put names or email addresses on spans. Renaming and + # hashing belong in the collector or gateway. # # contextKeys: # - {from: sub, to: user.id} # - thread_id - # - channel - # - # To derive a stable identifier without putting the original value on the - # span, set hash: hmac-sha256 and provide contextHashKeySecret: - # - {from: email, to: user.hash, hash: hmac-sha256} + # - gen_ai.conversation.id contextKeys: [] - # Secret providing KAGENT_TRACE_CONTEXT_HASH_KEY for hash: hmac-sha256 - # mappings. Injected into the controller and forwarded to agent runtimes. - # A Harness cannot set or replace it. - contextHashKeySecret: - name: "" - key: hmac-key exporter: otlp: endpoint: "" diff --git a/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py b/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py index 9ebe243eaf..6c3c611c9a 100644 --- a/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py +++ b/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py @@ -34,7 +34,7 @@ hitl_activated, now_timestamp, ) -from kagent.core.tracing import merge_caller_context_attributes +from kagent.core.tracing import detach_promoted_metadata, promote_message_metadata_to_baggage from kagent.core.tracing._span_processor import clear_kagent_span_attributes, set_kagent_span_attributes from pydantic import BaseModel @@ -120,6 +120,7 @@ async def execute(self, context: RequestContext, event_queue: EventQueue) -> Non runner: Runner | None = None context_token = None + promote_token = None try: self._translate_hitl_response(context) runner = await self._resolve_runner() @@ -132,10 +133,10 @@ async def execute(self, context: RequestContext, event_queue: EventQueue) -> Non "gen_ai.task.id": context.task_id, "gen_ai.conversation.id": run_request.session_id, } - # Allowlisted caller context joins the request-scoped bag rather - # than a single span, so tool, sub-agent, and model spans all - # carry it. Fill-if-absent so a caller cannot override kagent.user_id. - merge_caller_context_attributes(span_attributes, message=context.message) + # Allowlisted metadata is written into baggage so the baggage + # span processor stamps it on every span, including hops the + # runtime does not own. Fill-if-absent so existing baggage wins. + promote_token = promote_message_metadata_to_baggage(message=context.message) context_token = set_kagent_span_attributes( {key: value for key, value in span_attributes.items() if value is not None} ) @@ -182,6 +183,7 @@ async def execute(self, context: RequestContext, event_queue: EventQueue) -> Non finally: if context_token is not None: clear_kagent_span_attributes(context_token) + detach_promoted_metadata(promote_token) if runner is not None: await self._safe_close_runner(runner) diff --git a/python/packages/kagent-adk/tests/unittests/test_agent_executor.py b/python/packages/kagent-adk/tests/unittests/test_agent_executor.py index 393cf5f615..a315e588ab 100644 --- a/python/packages/kagent-adk/tests/unittests/test_agent_executor.py +++ b/python/packages/kagent-adk/tests/unittests/test_agent_executor.py @@ -107,7 +107,7 @@ async def execute(self, request_context, queue): @pytest.mark.asyncio -async def test_execute_merges_caller_context_without_eager_decode(monkeypatch): +async def test_execute_promotes_message_metadata_to_baggage(monkeypatch): context = _request_context() event_queue = object() runner = object() @@ -124,11 +124,11 @@ async def test_execute_merges_caller_context_without_eager_decode(monkeypatch): recorded = {} - def fake_merge(existing, metadata=None, context=None, message=None): - recorded["existing"] = dict(existing) + def fake_promote(*, metadata=None, context=None, message=None): recorded["message"] = message + return None - monkeypatch.setattr(executor_module, "merge_caller_context_attributes", fake_merge) + monkeypatch.setattr(executor_module, "promote_message_metadata_to_baggage", fake_promote) decode = MagicMock(return_value={"thread_id": "T1"}) if hasattr(executor_module, "read_message_metadata"): monkeypatch.setattr(executor_module, "read_message_metadata", decode) @@ -145,5 +145,4 @@ async def execute(self, request_context, queue): await executor.execute(context, event_queue) assert recorded["message"] is context.message - assert recorded["existing"]["kagent.user_id"] == "user-1" decode.assert_not_called() diff --git a/python/packages/kagent-core/pyproject.toml b/python/packages/kagent-core/pyproject.toml index 4aeb868817..d4c998e859 100644 --- a/python/packages/kagent-core/pyproject.toml +++ b/python/packages/kagent-core/pyproject.toml @@ -20,6 +20,7 @@ dependencies = [ "opentelemetry-instrumentation-httpx>=0.63b0,<0.64", "opentelemetry-instrumentation-fastapi>=0.63b0,<0.64", "opentelemetry-instrumentation-google-generativeai>=0.52.5,<0.53.0", + "opentelemetry-processor-baggage>=0.63b0,<0.64", "typing-extensions>=4.16.0", ] diff --git a/python/packages/kagent-core/src/kagent/core/tracing/__init__.py b/python/packages/kagent-core/src/kagent/core/tracing/__init__.py index 96d6cdcd40..194d317cd6 100644 --- a/python/packages/kagent-core/src/kagent/core/tracing/__init__.py +++ b/python/packages/kagent-core/src/kagent/core/tracing/__init__.py @@ -1,4 +1,16 @@ -from ._context_attributes import caller_context_attributes, merge_caller_context_attributes +from ._context_attributes import ( + allowed_baggage_key_predicate, + context_with_promoted_metadata, + detach_promoted_metadata, + promote_message_metadata_to_baggage, +) from ._utils import configure, force_flush -__all__ = ["caller_context_attributes", "configure", "force_flush", "merge_caller_context_attributes"] +__all__ = [ + "allowed_baggage_key_predicate", + "configure", + "context_with_promoted_metadata", + "detach_promoted_metadata", + "force_flush", + "promote_message_metadata_to_baggage", +] diff --git a/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py b/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py index ac90f8d2a9..391a31a2c3 100644 --- a/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py +++ b/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py @@ -1,14 +1,13 @@ -"""Promote allowlisted caller context onto every span of an agent request.""" +"""Promote allowlisted caller context into W3C baggage.""" from __future__ import annotations import functools -import hashlib -import hmac import json import os +from contextvars import Token from dataclasses import dataclass -from typing import Any, Optional +from typing import Any, Callable, Optional from a2a.types import Message from opentelemetry import baggage @@ -16,121 +15,94 @@ from kagent.core.a2a._consts import read_message_metadata -# Allowlist of caller-supplied context keys to promote onto agent spans. -# Unset or empty (the default) disables promotion entirely. +# Allowlist of caller-supplied context keys copied from W3C baggage onto +# agent spans. Unset or empty (the default) disables promotion entirely. # -# Accepts a comma-separated list of source keys, or a JSON array of strings -# and {from, to, hash} objects. +# Accepts a comma-separated list of keys, or a JSON array of strings and +# {from, to} objects. Destination names are used as-is. TRACE_CONTEXT_KEYS_ENV_VAR = "KAGENT_TRACE_CONTEXT_KEYS" -# HMAC key used when a mapping sets hash: hmac-sha256. Required for those -# entries; without it the hashed attribute is skipped rather than emitted -# in plaintext. -TRACE_CONTEXT_HASH_KEY_ENV_VAR = "KAGENT_TRACE_CONTEXT_HASH_KEY" - -# Namespaces custom promoted values so they cannot shadow a semantic -# convention attribute such as ``service.name``. Only the explicit registry -# names in ``_UNPREFIXED_ATTRIBUTES`` pass through as-is; see -# ``_span_attribute_name``. -CONTEXT_ATTRIBUTE_PREFIX = "kagent.context." - -HASH_HMAC_SHA256 = "hmac-sha256" - MAX_CONTEXT_KEYS = 32 MAX_CONTEXT_KEY_LENGTH = 64 -MAX_CONTEXT_VALUE_LENGTH = 256 - -# OpenTelemetry registry names callers may emit without a kagent.context. -# prefix. An invented name such as user.asdasd cannot masquerade as a -# semantic convention. -_UNPREFIXED_ATTRIBUTES = frozenset({"user.id", "user.hash", "enduser.id", "session.id"}) @dataclass(frozen=True) class _ContextMapping: source: str attribute: str - hash: str = "" -def caller_context_attributes( +def context_with_promoted_metadata( metadata: Optional[dict[str, Any]] = None, context: Optional[otel_context.Context] = None, message: Optional[Message] = None, -) -> dict[str, str]: - """Return the caller context values an operator allowlisted for tracing. - - Merge the result into the request-scoped attribute bag (see - ``set_kagent_span_attributes``) so every span of the request carries the - values: trace-level filtering in backends such as Langfuse matches on each - span, not only on the root. - - Values are read from W3C baggage first and then from the A2A message - metadata, which is the more specific source for a single message and - therefore wins — but only when the sanitised metadata scalar is non-empty, - so a blank metadata entry cannot wipe a baggage value. Both are untrusted - input, so keys must appear in the allowlist, values are stripped of control - characters and truncated, and attribute names go through - ``_span_attribute_name``. - - ``message`` is decoded only after the allowlist is known to be non-empty, - so the default (promotion off) does not pay for protobuf Struct conversion. - - Args: - metadata: A2A message metadata as a plain dict (may be ``None``). - context: OTel context to read baggage from. Defaults to the current one. - message: A2A message whose metadata is decoded lazily when ``metadata`` - is ``None`` and the allowlist is non-empty. - - Returns: - Attribute name to sanitised value. Empty when the allowlist is empty, - which is the default. +) -> otel_context.Context: + """Copy allowlisted A2A metadata into baggage when the destination is unset. + + Existing baggage members named in ``from`` are remapped onto ``to`` the + same way, so both sources share one path: the baggage span processor. + + ``message`` is decoded only after the allowlist is known to be non-empty. + Empty or non-scalar metadata does not wipe baggage. """ mappings = _allowed_context_mappings() + ctx = context if context is not None else otel_context.get_current() if not mappings: - return {} + return ctx if metadata is None and message is not None: metadata = read_message_metadata(message) + if metadata is None: + metadata = {} - bag = baggage.get_all(context) - attributes: dict[str, str] = {} + bag = baggage.get_all(ctx) for mapping in mappings: - value = _sanitize_context_value(bag.get(mapping.source)) - if metadata is not None: - scalar = _scalar_string(metadata.get(mapping.source)) - if scalar is not None: - cleaned = _sanitize_context_value(scalar) - if cleaned: - value = cleaned - if not value: + existing = bag.get(mapping.attribute) + if isinstance(existing, str) and existing.strip(): continue - if mapping.hash: - value = _hash_context_value(value, mapping.hash) - if not value: - continue - else: - value = _truncate_context_value(value) - name = _span_attribute_name(mapping.attribute) - if name in attributes: + value = "" + scalar = _scalar_string(metadata.get(mapping.source)) + if scalar is not None: + value = _sanitize_context_value(scalar) + if not value and mapping.source != mapping.attribute: + value = _sanitize_context_value(bag.get(mapping.source)) + if not value: continue - attributes[name] = value - return attributes + ctx = baggage.set_baggage(mapping.attribute, value, context=ctx) + bag = baggage.get_all(ctx) + return ctx -def merge_caller_context_attributes( - existing: dict[str, Any], +def promote_message_metadata_to_baggage( metadata: Optional[dict[str, Any]] = None, context: Optional[otel_context.Context] = None, message: Optional[Message] = None, -) -> None: - """Copy promoted caller context into *existing* without replacing keys. +) -> Optional[Token[otel_context.Context]]: + """Attach a context with allowlisted metadata written into baggage. - Caller-supplied data cannot override ``kagent.user_id``, ``kagent.app_name``, - or other attributes the runtime already stamped. + Returns a token that must be detached with + :func:`detach_promoted_metadata`, or ``None`` when nothing changed. """ - for key, value in caller_context_attributes(metadata, context, message=message).items(): - if key not in existing: - existing[key] = value + attached = otel_context.get_current() + source = context if context is not None else attached + promoted = context_with_promoted_metadata(metadata, source, message) + if promoted is attached: + return None + return otel_context.attach(promoted) + + +def detach_promoted_metadata(token: Optional[Token[otel_context.Context]]) -> None: + """Detach a token returned by :func:`promote_message_metadata_to_baggage`.""" + if token is not None: + otel_context.detach(token) + + +def allowed_baggage_key_predicate() -> Optional[Callable[[str], bool]]: + """Return a BaggageSpanProcessor predicate, or ``None`` when promotion is off.""" + mappings = _allowed_context_mappings() + if not mappings: + return None + allowed = frozenset(mapping.attribute for mapping in mappings) + return allowed.__contains__ @functools.cache @@ -138,8 +110,7 @@ def _allowed_context_mappings() -> tuple[_ContextMapping, ...]: """Parse the ``KAGENT_TRACE_CONTEXT_KEYS`` allowlist. Keys that are empty, over-long, or contain whitespace or control characters - are dropped, and the list is capped at ``MAX_CONTEXT_KEYS`` so a - misconfigured allowlist cannot inflate span cardinality without bound. + are dropped, and the list is capped at ``MAX_CONTEXT_KEYS``. Cached because the allowlist cannot change without a process restart. """ raw = os.getenv(TRACE_CONTEXT_KEYS_ENV_VAR, "").strip() @@ -153,7 +124,7 @@ def _allowed_context_mappings() -> tuple[_ContextMapping, ...]: def _parse_comma_allowlist(raw: str) -> list[_ContextMapping]: mappings: list[_ContextMapping] = [] for candidate in raw.split(","): - mapping = _new_context_mapping(candidate.strip(), "", "") + mapping = _new_context_mapping(candidate.strip(), "") if mapping is not None: mappings.append(mapping) return mappings @@ -169,17 +140,14 @@ def _parse_json_allowlist(raw: str) -> list[_ContextMapping]: mappings: list[_ContextMapping] = [] for item in items: if isinstance(item, str): - mapping = _new_context_mapping(item, "", "") + mapping = _new_context_mapping(item, "") elif isinstance(item, dict): from_key, from_ok = _json_string_field(item, "from") to_key, to_ok = _json_string_field(item, "to") - hash_alg, hash_ok = _json_string_field(item, "hash") - if not from_ok or not to_ok or not hash_ok: - # Non-string hash/to must drop the mapping. Coercing hash: 123 - # to "no hash" would emit the original value in plaintext. + if not from_ok or not to_ok: mapping = None else: - mapping = _new_context_mapping(from_key, to_key, hash_alg) + mapping = _new_context_mapping(from_key, to_key) else: mapping = None if mapping is not None: @@ -188,12 +156,7 @@ def _parse_json_allowlist(raw: str) -> list[_ContextMapping]: def _json_string_field(item: dict[str, Any], name: str) -> tuple[str, bool]: - """Return a JSON string field, or reject it when the type is invalid. - - Missing keys and JSON null are unset (empty string), matching Go's - encoding/json behaviour for string fields. Any other non-string type is - a malformed mapping and must fail closed. - """ + """Return a JSON string field, or reject it when the type is invalid.""" if name not in item: return "", True raw = item[name] @@ -204,26 +167,23 @@ def _json_string_field(item: dict[str, Any], name: str) -> tuple[str, bool]: return "", False -def _new_context_mapping(from_key: str, to_key: str, hash_alg: str) -> Optional[_ContextMapping]: +def _new_context_mapping(from_key: str, to_key: str) -> Optional[_ContextMapping]: from_key = from_key.strip() to_key = to_key.strip() - hash_alg = hash_alg.strip() if not from_key or len(from_key) > MAX_CONTEXT_KEY_LENGTH or not _is_attribute_key(from_key): return None if not to_key: to_key = from_key if len(to_key) > MAX_CONTEXT_KEY_LENGTH or not _is_attribute_key(to_key): return None - if hash_alg and hash_alg != HASH_HMAC_SHA256: - return None - return _ContextMapping(source=from_key, attribute=to_key, hash=hash_alg) + return _ContextMapping(source=from_key, attribute=to_key) def _cap_mappings(mappings: list[_ContextMapping]) -> list[_ContextMapping]: out: list[_ContextMapping] = [] - seen: set[tuple[str, str, str]] = set() + seen: set[tuple[str, str]] = set() for mapping in mappings: - identity = (mapping.source, mapping.attribute, mapping.hash) + identity = (mapping.source, mapping.attribute) if identity in seen: continue seen.add(identity) @@ -233,21 +193,7 @@ def _cap_mappings(mappings: list[_ContextMapping]) -> list[_ContextMapping]: return out -def _span_attribute_name(name: str) -> str: - """Return the name written onto the span. - - ``user.id``, ``user.hash``, ``enduser.id``, and ``session.id`` pass through - unprefixed so operators can use the semantic convention names. Everything - else is placed under ``kagent.context.`` so a caller-supplied - ``service.name`` or ``kagent.user_id`` cannot shadow the real one. - """ - if name in _UNPREFIXED_ATTRIBUTES: - return name - return CONTEXT_ATTRIBUTE_PREFIX + name - - def _is_attribute_key(key: str) -> bool: - """Report whether *key* is safe to use as a span attribute name.""" return not any(_is_control(char) or char.isspace() for char in key) @@ -258,12 +204,7 @@ def _is_control(char: str) -> bool: def _scalar_string(value: Any) -> Optional[str]: - """Render a JSON scalar from A2A message metadata as a string. - - Objects and arrays are skipped: they are unbounded in size and carry no - useful meaning as a span attribute value. ``None`` means "no scalar here", - which leaves any baggage value for the same key in place. - """ + """Render a JSON scalar from A2A message metadata as a string.""" if isinstance(value, bool): return "true" if value else "false" if isinstance(value, int): @@ -276,12 +217,7 @@ def _scalar_string(value: Any) -> Optional[str]: def _format_float(value: float) -> str: - """Match Go ``strconv.FormatFloat(v, 'f', -1, 64)``: never scientific notation. - - ``format(value, '.16f')`` rounds values such as ``1e-20`` to ``0``. Expand - Python's shortest unique repr from scientific form into fixed-point - instead, so tiny values keep their digits. - """ + """Match Go ``strconv.FormatFloat(v, 'f', -1, 64)``: never scientific notation.""" return _fixed_point_from_shortest(repr(value)) @@ -322,21 +258,3 @@ def _sanitize_context_value(value: Any) -> str: if not isinstance(value, str): return "" return "".join(char for char in value if not _is_control(char)).strip() - - -def _truncate_context_value(value: str) -> str: - return value[:MAX_CONTEXT_VALUE_LENGTH] - - -def _hash_context_value(value: str, algorithm: str) -> str: - """Hash *value* with the requested algorithm. - - Unknown algorithms and a missing HMAC key skip the attribute: never fall - back to putting the original value on the span. - """ - if algorithm != HASH_HMAC_SHA256: - return "" - key = os.getenv(TRACE_CONTEXT_HASH_KEY_ENV_VAR, "") - if not key: - return "" - return hmac.new(key.encode("utf-8"), value.encode("utf-8"), hashlib.sha256).hexdigest() diff --git a/python/packages/kagent-core/src/kagent/core/tracing/_utils.py b/python/packages/kagent-core/src/kagent/core/tracing/_utils.py index 2b2f78f791..add3b3b1a8 100644 --- a/python/packages/kagent-core/src/kagent/core/tracing/_utils.py +++ b/python/packages/kagent-core/src/kagent/core/tracing/_utils.py @@ -10,9 +10,11 @@ from opentelemetry.sdk._logs import LoggerProvider from opentelemetry.sdk._logs.export import BatchLogRecordProcessor from opentelemetry.sdk.resources import Resource +from opentelemetry.processor.baggage import BaggageSpanProcessor from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import BatchSpanProcessor +from ._context_attributes import allowed_baggage_key_predicate from ._span_processor import KagentAttributesSpanProcessor @@ -48,6 +50,19 @@ def _create_log_exporter(**kwargs): return OTLPLogExporter(**kwargs) +def _add_span_processors(provider: TracerProvider, otlp_processor) -> None: + """Install baggagecopy (if allowlisted), the OTLP exporter, then kagent attrs. + + Baggagecopy runs before the request-scoped bag so runtime attributes such as + gen_ai.conversation.id still win on conflict. + """ + provider.add_span_processor(otlp_processor) + predicate = allowed_baggage_key_predicate() + if predicate is not None: + provider.add_span_processor(BaggageSpanProcessor(predicate)) + provider.add_span_processor(KagentAttributesSpanProcessor()) + + def _resolve_otlp_timeout_seconds(signal: str) -> float: """ Resolve OTLP timeout env vars (milliseconds) into seconds for exporters. @@ -263,14 +278,12 @@ def configure( current_provider = trace.get_tracer_provider() if isinstance(current_provider, TracerProvider): # TracerProvider already exists, just add our processors to it - current_provider.add_span_processor(processor) - current_provider.add_span_processor(KagentAttributesSpanProcessor()) + _add_span_processors(current_provider, processor) logging.info("Added OTLP processors to existing TracerProvider") else: # No provider set, create new one tracer_provider = TracerProvider(resource=resource) - tracer_provider.add_span_processor(processor) - tracer_provider.add_span_processor(KagentAttributesSpanProcessor()) + _add_span_processors(tracer_provider, processor) trace.set_tracer_provider(tracer_provider) logging.info("Created new TracerProvider") diff --git a/python/packages/kagent-core/tests/test_caller_context_attributes.py b/python/packages/kagent-core/tests/test_caller_context_attributes.py index 25ec635d67..e2637c35cf 100644 --- a/python/packages/kagent-core/tests/test_caller_context_attributes.py +++ b/python/packages/kagent-core/tests/test_caller_context_attributes.py @@ -1,21 +1,21 @@ -import hashlib -import hmac -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest from opentelemetry import baggage from opentelemetry import context as otel_context +from opentelemetry.processor.baggage import BaggageSpanProcessor from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter -from kagent.core.tracing import caller_context_attributes, merge_caller_context_attributes +from kagent.core.tracing import ( + allowed_baggage_key_predicate, + context_with_promoted_metadata, + detach_promoted_metadata, + promote_message_metadata_to_baggage, +) from kagent.core.tracing._context_attributes import ( - CONTEXT_ATTRIBUTE_PREFIX, - MAX_CONTEXT_KEY_LENGTH, MAX_CONTEXT_KEYS, - MAX_CONTEXT_VALUE_LENGTH, - TRACE_CONTEXT_HASH_KEY_ENV_VAR, TRACE_CONTEXT_KEYS_ENV_VAR, _allowed_context_mappings, ) @@ -26,6 +26,14 @@ ) +def set_allowlist(monkeypatch, allowlist: str) -> None: + if allowlist: + monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, allowlist) + else: + monkeypatch.delenv(TRACE_CONTEXT_KEYS_ENV_VAR, raising=False) + _allowed_context_mappings.cache_clear() + + def baggage_context(members: dict[str, str]) -> otel_context.Context: context = otel_context.Context() for key, value in members.items(): @@ -33,318 +41,191 @@ def baggage_context(members: dict[str, str]) -> otel_context.Context: return context -def hmac_sha256_hex(key: str, value: str) -> str: - return hmac.new(key.encode("utf-8"), value.encode("utf-8"), hashlib.sha256).hexdigest() +def baggage_values(context: otel_context.Context) -> dict[str, str]: + return {key: value for key, value in baggage.get_all(context).items()} -class TestCallerContextAttributes: - """Tests for allowlist-driven promotion of caller context onto spans.""" +@pytest.fixture(autouse=True) +def _clear_allowlist_cache(): + _allowed_context_mappings.cache_clear() + yield + _allowed_context_mappings.cache_clear() - def test_empty_allowlist_disables_promotion(self, monkeypatch): - monkeypatch.delenv(TRACE_CONTEXT_KEYS_ENV_VAR, raising=False) - assert ( - caller_context_attributes( - {"thread_id": "T1"}, - baggage_context({"sub": "opaque-subject"}), - ) - == {} - ) - def test_promotes_allowlisted_baggage(self, monkeypatch): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "sub,thread_id") - assert caller_context_attributes(None, baggage_context({"sub": "opaque-subject", "thread_id": "T123"})) == { - "kagent.context.sub": "opaque-subject", - "kagent.context.thread_id": "T123", +class TestContextWithPromotedMetadata: + def test_empty_allowlist_is_a_noop(self, monkeypatch): + set_allowlist(monkeypatch, "") + ctx = baggage_context({"user.id": "opaque-subject"}) + assert baggage_values(context_with_promoted_metadata({"user.id": "from-metadata"}, ctx)) == { + "user.id": "opaque-subject" + } + + def test_leaves_existing_baggage_in_place(self, monkeypatch): + set_allowlist(monkeypatch, "user.id") + ctx = baggage_context({"user.id": "from-baggage"}) + assert baggage_values(context_with_promoted_metadata({"user.id": "from-metadata"}, ctx)) == { + "user.id": "from-baggage" + } + + def test_promotes_allowlisted_metadata(self, monkeypatch): + set_allowlist(monkeypatch, "user.id,thread_id") + assert baggage_values( + context_with_promoted_metadata({"user.id": "opaque-subject", "thread_id": "T123", "secret": "nope"}) + ) == {"user.id": "opaque-subject", "thread_id": "T123"} + + def test_empty_metadata_does_not_wipe_baggage(self, monkeypatch): + set_allowlist(monkeypatch, "user.id") + ctx = baggage_context({"user.id": "from-baggage"}) + assert baggage_values(context_with_promoted_metadata({"user.id": " "}, ctx)) == {"user.id": "from-baggage"} + + def test_remaps_baggage_from_onto_to(self, monkeypatch): + set_allowlist(monkeypatch, '[{"from":"sub","to":"user.id"}]') + ctx = baggage_context({"sub": "opaque-subject"}) + assert baggage_values(context_with_promoted_metadata({}, ctx)) == { + "sub": "opaque-subject", + "user.id": "opaque-subject", } - def test_promotes_allowlisted_message_metadata(self, monkeypatch): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "thread_id,channel") - assert caller_context_attributes({"thread_id": "1717171.4242", "channel": "C0AB1"}) == { - "kagent.context.thread_id": "1717171.4242", - "kagent.context.channel": "C0AB1", + def test_does_not_remap_over_existing_to(self, monkeypatch): + set_allowlist(monkeypatch, '[{"from":"sub","to":"user.id"}]') + ctx = baggage_context({"sub": "from-sub", "user.id": "already"}) + assert baggage_values(context_with_promoted_metadata({}, ctx)) == { + "sub": "from-sub", + "user.id": "already", } - def test_message_metadata_overrides_baggage(self, monkeypatch): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "sub") - assert caller_context_attributes( - {"sub": "from-metadata"}, - baggage_context({"sub": "from-baggage"}), - ) == {"kagent.context.sub": "from-metadata"} - - def test_empty_metadata_does_not_override_baggage(self, monkeypatch): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "sub") - assert caller_context_attributes( - {"sub": " "}, - baggage_context({"sub": "from-baggage"}), - ) == {"kagent.context.sub": "from-baggage"} - - def test_ignores_keys_outside_the_allowlist(self, monkeypatch): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "thread_id") - assert caller_context_attributes( - {"thread_id": "T1", "extra": "nope"}, - baggage_context({"secret.token": "s3cret"}), - ) == {"kagent.context.thread_id": "T1"} + def test_maps_metadata_from_onto_to(self, monkeypatch): + set_allowlist(monkeypatch, '[{"from":"sub","to":"user.id"},"gen_ai.conversation.id"]') + assert baggage_values( + context_with_promoted_metadata({"sub": "opaque-subject", "gen_ai.conversation.id": "sess-1"}) + ) == {"user.id": "opaque-subject", "gen_ai.conversation.id": "sess-1"} + + def test_skips_non_scalars(self, monkeypatch): + set_allowlist(monkeypatch, "value") + assert baggage_values(context_with_promoted_metadata({"value": {"nested": True}})) == {} @pytest.mark.parametrize( ("value", "expected"), [ (True, "true"), (False, "false"), - (7, "7"), - # protobuf Struct has one numeric type, so JSON integers arrive as floats. - (3.0, "3"), - (2.5, "2.5"), - (1e20, "100000000000000000000"), - (1e-20, "0.00000000000000000001"), + (3, "3"), + (1.5, "1.5"), ], ) - def test_renders_scalar_metadata_types(self, monkeypatch, value, expected): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "value") - assert caller_context_attributes({"value": value}) == {"kagent.context.value": expected} - - @pytest.mark.parametrize("value", [{"a": "b"}, ["a"], "", None]) - def test_skips_non_scalar_and_empty_metadata_values(self, monkeypatch, value): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "value") - assert caller_context_attributes({"value": value}) == {} + def test_renders_scalar_metadata(self, monkeypatch, value, expected): + set_allowlist(monkeypatch, "value") + assert baggage_values(context_with_promoted_metadata({"value": value})) == {"value": expected} def test_strips_control_characters(self, monkeypatch): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "note") - assert caller_context_attributes({"note": "line\nbreak\tand\x00nul"}) == { - "kagent.context.note": "linebreakandnul" + set_allowlist(monkeypatch, "note") + assert baggage_values(context_with_promoted_metadata({"note": "line\nbreak\tand\x00nul"})) == { + "note": "linebreakandnul" } - def test_truncates_long_values(self, monkeypatch): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "note") - promoted = caller_context_attributes({"note": "a" * (MAX_CONTEXT_VALUE_LENGTH * 2)}) - assert len(promoted["kagent.context.note"]) == MAX_CONTEXT_VALUE_LENGTH - - def test_cannot_shadow_semantic_conventions(self, monkeypatch): - """Custom keys still cannot replace service.name. Registry names are the exception.""" - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "service.name") - promoted = caller_context_attributes({"service.name": "impostor"}) - assert "service.name" not in promoted - assert promoted == {f"{CONTEXT_ATTRIBUTE_PREFIX}service.name": "impostor"} - - def test_registry_attributes_stay_unprefixed(self, monkeypatch): - """Only the explicit registry set is left unprefixed.""" - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "user.id,user.hash,enduser.id,session.id,channel") - assert caller_context_attributes( - { - "user.id": "opaque-subject", - "user.hash": "abc", - "enduser.id": "end-user", - "session.id": "sess-1", - "channel": "C0AB1", - } - ) == { - "user.id": "opaque-subject", - "user.hash": "abc", - "enduser.id": "end-user", - "session.id": "sess-1", - "kagent.context.channel": "C0AB1", - } - - def test_unknown_user_attributes_are_prefixed(self, monkeypatch): - monkeypatch.setenv( - TRACE_CONTEXT_KEYS_ENV_VAR, - '[{"from":"x","to":"user.asdasd"},{"from":"y","to":"user.email"}]', - ) - assert caller_context_attributes({"x": "nope", "y": "also-nope"}) == { - "kagent.context.user.asdasd": "nope", - "kagent.context.user.email": "also-nope", - } - - def test_session_id_is_unprefixed_but_session_foo_is_not(self, monkeypatch): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "session.id,session.foo") - assert caller_context_attributes({"session.id": "sess-1", "session.foo": "other"}) == { - "session.id": "sess-1", - "kagent.context.session.foo": "other", - } - - def test_maps_source_keys_onto_registry_names(self, monkeypatch): - monkeypatch.setenv( - TRACE_CONTEXT_KEYS_ENV_VAR, - '[{"from":"sub","to":"user.id"},{"from":"thread_id","to":"session.id"},"channel"]', - ) - assert caller_context_attributes({"sub": "opaque-subject", "thread_id": "T123", "channel": "C0AB1"}) == { - "user.id": "opaque-subject", - "session.id": "T123", - "kagent.context.channel": "C0AB1", - } - - def test_kagent_destination_names_are_prefixed(self, monkeypatch): - monkeypatch.setenv( - TRACE_CONTEXT_KEYS_ENV_VAR, - '[{"from":"uid","to":"kagent.user_id"},{"from":"tid","to":"kagent.thread_id"}]', - ) - assert caller_context_attributes({"uid": "attacker", "tid": "T123"}) == { - "kagent.context.kagent.user_id": "attacker", - "kagent.context.kagent.thread_id": "T123", - } + def test_drops_keys_with_whitespace(self, monkeypatch): + set_allowlist(monkeypatch, "bad key,user.id") + assert baggage_values(context_with_promoted_metadata({"bad key": "x", "user.id": "ok"})) == {"user.id": "ok"} - def test_invalid_json_allowlist_promotes_nothing(self, monkeypatch): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, '[{"from":"sub"') - assert caller_context_attributes({"sub": "opaque-subject"}) == {} - - def test_hashes_with_hmac_sha256(self, monkeypatch): - key = "test-hmac-key" - monkeypatch.setenv( - TRACE_CONTEXT_KEYS_ENV_VAR, - '[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]', - ) - monkeypatch.setenv(TRACE_CONTEXT_HASH_KEY_ENV_VAR, key) - promoted = caller_context_attributes({"email": "ada@example.com"}) - assert promoted == {"user.hash": hmac_sha256_hex(key, "ada@example.com")} - assert not any("@example.com" in value for value in promoted.values()) - - def test_hash_without_key_emits_nothing(self, monkeypatch): - """Missing HMAC key must not fall back to putting the original value on the span.""" - monkeypatch.setenv( - TRACE_CONTEXT_KEYS_ENV_VAR, - '[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]', - ) - monkeypatch.delenv(TRACE_CONTEXT_HASH_KEY_ENV_VAR, raising=False) - assert caller_context_attributes({"email": "ada@example.com"}) == {} - - def test_unknown_hash_emits_nothing(self, monkeypatch): - monkeypatch.setenv( - TRACE_CONTEXT_KEYS_ENV_VAR, - '[{"from":"email","to":"user.hash","hash":"md5"}]', - ) - monkeypatch.setenv(TRACE_CONTEXT_HASH_KEY_ENV_VAR, "test-hmac-key") - assert caller_context_attributes({"email": "ada@example.com"}) == {} - - def test_non_string_hash_does_not_emit_plaintext(self, monkeypatch): - """A numeric hash field must drop the mapping, not skip hashing.""" - monkeypatch.setenv( - TRACE_CONTEXT_KEYS_ENV_VAR, - '[{"from":"email","to":"user.hash","hash":123}]', - ) - monkeypatch.setenv(TRACE_CONTEXT_HASH_KEY_ENV_VAR, "test-hmac-key") - promoted = caller_context_attributes({"email": "ada@example.com"}) - assert promoted == {} - assert not any("ada@example.com" in value for value in promoted.values()) - - def test_non_string_to_drops_the_mapping(self, monkeypatch): - monkeypatch.setenv( - TRACE_CONTEXT_KEYS_ENV_VAR, - '[{"from":"thread_id","to":123}]', - ) - assert caller_context_attributes({"thread_id": "T1"}) == {} + def test_caps_the_allowlist(self, monkeypatch): + keys = [f"k{i}" for i in range(MAX_CONTEXT_KEYS + 4)] + set_allowlist(monkeypatch, ",".join(keys)) + got = baggage_values(context_with_promoted_metadata({key: "v" for key in keys})) + assert len(got) == MAX_CONTEXT_KEYS def test_skips_message_decode_when_allowlist_empty(self, monkeypatch): - monkeypatch.delenv(TRACE_CONTEXT_KEYS_ENV_VAR, raising=False) - with patch( - "kagent.core.tracing._context_attributes.read_message_metadata", - return_value={"thread_id": "T1"}, - ) as decode: - assert caller_context_attributes(message=MagicMock()) == {} - decode.assert_not_called() - - def test_decodes_message_when_allowlist_is_set(self, monkeypatch): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "thread_id") + set_allowlist(monkeypatch, "") + decode = MagicMock(return_value={"thread_id": "T1"}) + monkeypatch.setattr("kagent.core.tracing._context_attributes.read_message_metadata", decode) + assert baggage_values(context_with_promoted_metadata(message=MagicMock())) == {} + decode.assert_not_called() + + def test_decodes_message_after_allowlist_check(self, monkeypatch): + set_allowlist(monkeypatch, "thread_id") message = MagicMock() - with patch( + monkeypatch.setattr( "kagent.core.tracing._context_attributes.read_message_metadata", - return_value={"thread_id": "T1"}, - ) as decode: - assert caller_context_attributes(message=message) == {"kagent.context.thread_id": "T1"} - decode.assert_called_once_with(message) - - -class TestMergeCallerContextAttributes: - def test_does_not_override_existing_keys(self, monkeypatch): - monkeypatch.setenv( - TRACE_CONTEXT_KEYS_ENV_VAR, - '[{"from":"sub","to":"user.id"},"thread_id"]', - ) - existing = {"user.id": "runtime-user", "kagent.user_id": "runtime-user"} - merge_caller_context_attributes( - existing, - {"sub": "attacker", "thread_id": "T1"}, + lambda _message: {"thread_id": "T1"}, ) - assert existing == { - "user.id": "runtime-user", - "kagent.user_id": "runtime-user", - "kagent.context.thread_id": "T1", - } - + assert baggage_values(context_with_promoted_metadata(message=message)) == {"thread_id": "T1"} -class TestAllowedContextMappings: - """Tests for allowlist parsing and its bounds.""" - def test_caps_list_length(self, monkeypatch): - keys = ",".join(f"key{index}" for index in range(MAX_CONTEXT_KEYS * 2)) - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, keys) - assert len(_allowed_context_mappings()) == MAX_CONTEXT_KEYS +def test_allowed_baggage_key_predicate(monkeypatch): + set_allowlist(monkeypatch, "") + assert allowed_baggage_key_predicate() is None - def test_drops_over_long_and_duplicate_keys(self, monkeypatch): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, f"a,a,{'b' * (MAX_CONTEXT_KEY_LENGTH + 1)},c") - assert [mapping.source for mapping in _allowed_context_mappings()] == ["a", "c"] + set_allowlist(monkeypatch, '[{"from":"sub","to":"user.id"},"thread_id"]') + predicate = allowed_baggage_key_predicate() + assert predicate is not None + assert predicate("user.id") + assert predicate("thread_id") + assert not predicate("sub") + assert not predicate("secret") - def test_counts_key_length_in_code_points(self, monkeypatch): - kept = "键" * 40 - dropped = "键" * 65 - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, f"{kept},{dropped}") - assert [mapping.source for mapping in _allowed_context_mappings()] == [kept] - def test_drops_keys_that_are_not_valid_attribute_names(self, monkeypatch): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "good, bad key ,\tanother\tbad") - assert [mapping.source for mapping in _allowed_context_mappings()] == ["good"] +def test_promote_attaches_and_detaches(monkeypatch): + set_allowlist(monkeypatch, "user.id") + token = promote_message_metadata_to_baggage(metadata={"user.id": "opaque-subject"}) + assert baggage.get_baggage("user.id") == "opaque-subject" + detach_promoted_metadata(token) + assert baggage.get_baggage("user.id") is None - def test_empty_allowlist_disables_promotion(self, monkeypatch): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, " , ,") - assert _allowed_context_mappings() == () - def test_cache_survives_env_change_until_cleared(self, monkeypatch): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "a") - assert [mapping.source for mapping in _allowed_context_mappings()] == ["a"] - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, "b") - assert [mapping.source for mapping in _allowed_context_mappings()] == ["a"] - _allowed_context_mappings.cache_clear() - assert [mapping.source for mapping in _allowed_context_mappings()] == ["b"] +def _span_attrs(spans, name: str) -> dict: + for span in spans: + if span.name == name: + return dict(span.attributes or {}) + raise AssertionError(f"span {name} not found") -class TestPromotedAttributesReachEverySpan: - """Langfuse and comparable backends filter on attributes present on each - span, so promoted values must reach descendants, not just the root span.""" +def test_caller_context_lands_on_every_span(monkeypatch): + set_allowlist(monkeypatch, "user.id") + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + provider.add_span_processor(BaggageSpanProcessor(allowed_baggage_key_predicate())) + provider.add_span_processor(KagentAttributesSpanProcessor()) - @staticmethod - def record_spans(span_attributes: dict) -> dict[str, dict]: - exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - provider.add_span_processor(KagentAttributesSpanProcessor()) + token = promote_message_metadata_to_baggage( + metadata=None, + context=baggage_context({"user.id": "opaque-subject"}), + ) + attrs_token = set_kagent_span_attributes({"kagent.user_id": "runtime-user"}) + try: tracer = provider.get_tracer("test") - - token = set_kagent_span_attributes(span_attributes) - try: - with tracer.start_as_current_span("root"): - with tracer.start_as_current_span("execute_tool"): - with tracer.start_as_current_span("generate_content"): - pass - finally: - clear_kagent_span_attributes(token) - provider.shutdown() - - return {span.name: dict(span.attributes or {}) for span in exporter.get_finished_spans()} - - def test_flag_on_stamps_every_span(self, monkeypatch): - monkeypatch.setenv(TRACE_CONTEXT_KEYS_ENV_VAR, '[{"from":"sub","to":"user.id"}]') - promoted = caller_context_attributes(None, baggage_context({"sub": "opaque-subject"})) - - spans = self.record_spans(promoted) - - assert set(spans) == {"root", "execute_tool", "generate_content"} - for attributes in spans.values(): - assert attributes["user.id"] == "opaque-subject" - - def test_flag_off_leaves_spans_unchanged(self, monkeypatch): - monkeypatch.delenv(TRACE_CONTEXT_KEYS_ENV_VAR, raising=False) - promoted = caller_context_attributes({"thread_id": "T1"}, baggage_context({"sub": "opaque-subject"})) - - spans = self.record_spans(promoted) - - for attributes in spans.values(): - assert "user.id" not in attributes - assert not [key for key in attributes if key.startswith(CONTEXT_ATTRIBUTE_PREFIX)] + with tracer.start_as_current_span("invocation"): + with tracer.start_as_current_span("tool"): + with tracer.start_as_current_span("generate_content"): + pass + finally: + clear_kagent_span_attributes(attrs_token) + detach_promoted_metadata(token) + provider.shutdown() + + for name in ("invocation", "tool", "generate_content"): + attrs = _span_attrs(exporter.get_finished_spans(), name) + assert attrs["user.id"] == "opaque-subject" + assert attrs["kagent.user_id"] == "runtime-user" + + +def test_runtime_attributes_win_over_allowlisted_baggage(monkeypatch): + set_allowlist(monkeypatch, "gen_ai.conversation.id") + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + provider.add_span_processor(BaggageSpanProcessor(allowed_baggage_key_predicate())) + provider.add_span_processor(KagentAttributesSpanProcessor()) + + token = otel_context.attach(baggage_context({"gen_ai.conversation.id": "from-caller"})) + attrs_token = set_kagent_span_attributes({"gen_ai.conversation.id": "from-runtime"}) + try: + tracer = provider.get_tracer("test") + with tracer.start_as_current_span("generate_content"): + pass + finally: + clear_kagent_span_attributes(attrs_token) + otel_context.detach(token) + provider.shutdown() + + attrs = _span_attrs(exporter.get_finished_spans(), "generate_content") + assert attrs["gen_ai.conversation.id"] == "from-runtime" diff --git a/python/packages/kagent-core/tests/test_tracing_configure.py b/python/packages/kagent-core/tests/test_tracing_configure.py index b1415c66d8..7d1fe78a51 100644 --- a/python/packages/kagent-core/tests/test_tracing_configure.py +++ b/python/packages/kagent-core/tests/test_tracing_configure.py @@ -204,7 +204,7 @@ def test_otel_sdk_default_propagator_includes_baggage(): """The OTEL SDK must propagate W3C Baggage by default. Baggage is how caller identity and context reach an agent and its - sub-agents (see caller_context_attributes). If an OTEL SDK upgrade drops + sub-agents. If an OTEL SDK upgrade drops baggage from the default propagator, this test will fail and explicit configuration will be needed. """ diff --git a/python/packages/kagent-crewai/src/kagent/crewai/_executor.py b/python/packages/kagent-crewai/src/kagent/crewai/_executor.py index c99ff97e82..013fcf3698 100644 --- a/python/packages/kagent-crewai/src/kagent/crewai/_executor.py +++ b/python/packages/kagent-crewai/src/kagent/crewai/_executor.py @@ -23,7 +23,7 @@ ) from google.protobuf.json_format import MessageToDict from kagent.core.a2a import get_kagent_metadata_key, now_timestamp -from kagent.core.tracing import merge_caller_context_attributes +from kagent.core.tracing import detach_promoted_metadata, promote_message_metadata_to_baggage from kagent.core.tracing._span_processor import ( clear_kagent_span_attributes, set_kagent_span_attributes, @@ -70,6 +70,7 @@ async def execute( # Convert the a2a request to kagent span attributes. span_attributes = _convert_a2a_request_to_span_attributes(context) + promote_token = promote_message_metadata_to_baggage(message=context.message) # Set kagent span attributes for all spans in context. context_token = set_kagent_span_attributes(span_attributes) try: @@ -169,6 +170,7 @@ async def execute( ) finally: clear_kagent_span_attributes(context_token) + detach_promoted_metadata(promote_token) def _get_user_id(request: RequestContext) -> str: @@ -194,9 +196,4 @@ def _convert_a2a_request_to_span_attributes( if request.task_id: span_attributes["gen_ai.task.id"] = request.task_id - # Allowlisted caller context joins the request-scoped bag rather than a - # single span, so tool, sub-agent, and model spans all carry it. - # Fill-if-absent so a caller cannot override kagent.user_id. - merge_caller_context_attributes(span_attributes, message=request.message) - return span_attributes diff --git a/python/packages/kagent-crewai/tests/test_executor.py b/python/packages/kagent-crewai/tests/test_executor.py index a59751fea9..9b3b49c0b4 100644 --- a/python/packages/kagent-crewai/tests/test_executor.py +++ b/python/packages/kagent-crewai/tests/test_executor.py @@ -74,34 +74,9 @@ async def test_execute_falls_back_to_text_input_without_datapart(): _assert_content_artifact_closes_stream(events) -def test_span_attributes_skip_decode_when_allowlist_empty(monkeypatch): - from kagent.core.tracing._context_attributes import TRACE_CONTEXT_KEYS_ENV_VAR, _allowed_context_mappings - - monkeypatch.delenv(TRACE_CONTEXT_KEYS_ENV_VAR, raising=False) - _allowed_context_mappings.cache_clear() - decode = MagicMock(return_value={"thread_id": "T1"}) - monkeypatch.setattr("kagent.core.tracing._context_attributes.read_message_metadata", decode) - import kagent.crewai._executor as executor_module - - if hasattr(executor_module, "read_message_metadata"): - monkeypatch.setattr(executor_module, "read_message_metadata", decode) - +def test_span_attributes_are_runtime_only(): attrs = _convert_a2a_request_to_span_attributes(_request_context(Part(text="hello"))) - decode.assert_not_called() - assert "kagent.context.thread_id" not in attrs assert attrs["kagent.user_id"] - - -def test_span_attributes_do_not_override_runtime_user_id(monkeypatch): - fake = MagicMock(return_value={"kagent.user_id": "attacker", "kagent.context.thread_id": "T1"}) - monkeypatch.setattr("kagent.core.tracing._context_attributes.caller_context_attributes", fake) - import kagent.crewai._executor as executor_module - - if hasattr(executor_module, "caller_context_attributes"): - monkeypatch.setattr(executor_module, "caller_context_attributes", fake) - - attrs = _convert_a2a_request_to_span_attributes(_request_context(Part(text="hello"))) - - assert attrs["kagent.user_id"] != "attacker" - assert attrs["kagent.context.thread_id"] == "T1" + assert "kagent.context.thread_id" not in attrs + assert "user.id" not in attrs diff --git a/python/packages/kagent-langgraph/src/kagent/langgraph/_executor.py b/python/packages/kagent-langgraph/src/kagent/langgraph/_executor.py index 76f94aa42e..9d1b191d80 100644 --- a/python/packages/kagent-langgraph/src/kagent/langgraph/_executor.py +++ b/python/packages/kagent-langgraph/src/kagent/langgraph/_executor.py @@ -43,7 +43,7 @@ require_ask_user_response, require_tool_approval_response, ) -from kagent.core.tracing import merge_caller_context_attributes +from kagent.core.tracing import detach_promoted_metadata, promote_message_metadata_to_baggage from kagent.core.tracing._span_processor import ( clear_kagent_span_attributes, set_kagent_span_attributes, @@ -397,6 +397,7 @@ async def execute( # Convert the a2a request to kagent span attributes. span_attributes = _convert_a2a_request_to_span_attributes(context) + promote_token = promote_message_metadata_to_baggage(message=context.message) # Set kagent span attributes for all spans in context. context_token = set_kagent_span_attributes(span_attributes) try: @@ -505,6 +506,7 @@ async def execute( ) finally: clear_kagent_span_attributes(context_token) + detach_promoted_metadata(promote_token) def _get_user_id(request: RequestContext) -> str: @@ -542,9 +544,4 @@ def _convert_a2a_request_to_span_attributes( if request.task_id: span_attributes["gen_ai.task.id"] = request.task_id - # Allowlisted caller context joins the request-scoped bag rather than a - # single span, so tool, sub-agent, and model spans all carry it. - # Fill-if-absent so a caller cannot override kagent.user_id. - merge_caller_context_attributes(span_attributes, message=request.message) - return span_attributes diff --git a/python/packages/kagent-openai/src/kagent/openai/_agent_executor.py b/python/packages/kagent-openai/src/kagent/openai/_agent_executor.py index b8c9321ce2..b3e6e20fe3 100644 --- a/python/packages/kagent-openai/src/kagent/openai/_agent_executor.py +++ b/python/packages/kagent-openai/src/kagent/openai/_agent_executor.py @@ -36,7 +36,7 @@ from agents.memory.session import SessionABC from agents.run import Runner from kagent.core.a2a import get_kagent_metadata_key, now_timestamp -from kagent.core.tracing import merge_caller_context_attributes +from kagent.core.tracing import detach_promoted_metadata, promote_message_metadata_to_baggage from kagent.core.tracing._span_processor import ( clear_kagent_span_attributes, set_kagent_span_attributes, @@ -180,11 +180,13 @@ async def execute( raise ValueError("A2A request must have a message") span_attributes = _convert_a2a_request_to_span_attributes(context) + promote_token = promote_message_metadata_to_baggage(message=context.message) context_token = set_kagent_span_attributes(span_attributes) try: await self._execute_agent(context, event_queue) finally: clear_kagent_span_attributes(context_token) + detach_promoted_metadata(promote_token) async def _execute_agent( self, @@ -316,9 +318,4 @@ def _convert_a2a_request_to_span_attributes( if request.task_id: span_attributes["gen_ai.task.id"] = request.task_id - # Allowlisted caller context joins the request-scoped bag rather than a - # single span, so tool, sub-agent, and model spans all carry it. - # Fill-if-absent so a caller cannot override kagent.user_id. - merge_caller_context_attributes(span_attributes, message=request.message) - return span_attributes diff --git a/python/packages/kagent-openai/tests/test_caller_context.py b/python/packages/kagent-openai/tests/test_caller_context.py index 5254523f71..5a564982ff 100644 --- a/python/packages/kagent-openai/tests/test_caller_context.py +++ b/python/packages/kagent-openai/tests/test_caller_context.py @@ -16,34 +16,10 @@ def _request_context() -> RequestContext: ) -def test_span_attributes_skip_decode_when_allowlist_empty(monkeypatch): - from kagent.core.tracing._context_attributes import TRACE_CONTEXT_KEYS_ENV_VAR, _allowed_context_mappings - - monkeypatch.delenv(TRACE_CONTEXT_KEYS_ENV_VAR, raising=False) - _allowed_context_mappings.cache_clear() - decode = MagicMock(return_value={"thread_id": "T1"}) - monkeypatch.setattr("kagent.core.tracing._context_attributes.read_message_metadata", decode) - import kagent.openai._agent_executor as executor_module - - if hasattr(executor_module, "read_message_metadata"): - monkeypatch.setattr(executor_module, "read_message_metadata", decode) - +def test_span_attributes_are_runtime_only(): attrs = _convert_a2a_request_to_span_attributes(_request_context()) - decode.assert_not_called() - assert "kagent.context.thread_id" not in attrs assert "kagent.user_id" in attrs - - -def test_span_attributes_do_not_override_runtime_user_id(monkeypatch): - fake = MagicMock(return_value={"kagent.user_id": "attacker", "kagent.context.thread_id": "T1"}) - monkeypatch.setattr("kagent.core.tracing._context_attributes.caller_context_attributes", fake) - import kagent.openai._agent_executor as executor_module - - if hasattr(executor_module, "caller_context_attributes"): - monkeypatch.setattr(executor_module, "caller_context_attributes", fake) - - attrs = _convert_a2a_request_to_span_attributes(_request_context()) - - assert attrs["kagent.user_id"] != "attacker" - assert attrs["kagent.context.thread_id"] == "T1" + assert attrs["gen_ai.conversation.id"] == "ctx-1" + assert "kagent.context.thread_id" not in attrs + assert "user.id" not in attrs diff --git a/python/uv.lock b/python/uv.lock index 7e3578d9e0..e4a6b7fedf 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -250,9 +250,9 @@ name = "aiologic" version = "0.17.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "sniffio" }, - { name = "typing-extensions" }, - { name = "wrapt" }, + { name = "sniffio", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "wrapt", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/f1/7a/d51f2fde1e8ae8a83431f8e97b7a71e9358cdb1d4d2ce6be387fa44d68de/aiologic-0.17.1.tar.gz", hash = "sha256:2e1b93b9e88ced318c2a63ad7b382688f40cbfe40e3d42258d49dc9c5aea179d", size = 252354, upload-time = "2026-06-27T20:41:33.25Z" } wheels = [ @@ -872,8 +872,8 @@ name = "culsans" version = "0.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "aiologic" }, - { name = "typing-extensions" }, + { name = "aiologic", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/d9/e3/49afa1bc180e0d28008ec6bcdf82a4072d1c7a41032b5b759b60814ca4b0/culsans-0.11.0.tar.gz", hash = "sha256:0b43d0d05dce6106293d114c86e3fb4bfc63088cfe8ff08ed3fe36891447fe33", size = 107546, upload-time = "2025-12-31T23:15:38.196Z" } wheels = [ @@ -1885,6 +1885,7 @@ dependencies = [ { name = "opentelemetry-instrumentation-google-generativeai" }, { name = "opentelemetry-instrumentation-httpx" }, { name = "opentelemetry-instrumentation-openai" }, + { name = "opentelemetry-processor-baggage" }, { name = "opentelemetry-sdk" }, { name = "typing-extensions" }, ] @@ -1901,6 +1902,7 @@ requires-dist = [ { name = "opentelemetry-instrumentation-google-generativeai", specifier = ">=0.52.5,<0.53.0" }, { name = "opentelemetry-instrumentation-httpx", specifier = ">=0.63b0,<0.64" }, { name = "opentelemetry-instrumentation-openai", specifier = ">=0.52.5,<0.53.0" }, + { name = "opentelemetry-processor-baggage", specifier = ">=0.63b0,<0.64" }, { name = "opentelemetry-sdk", specifier = ">=1.42.1,<1.43.0" }, { name = "typing-extensions", specifier = ">=4.16.0" }, ] @@ -3031,6 +3033,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5a/17/b3054b866d18d49997ba9d7114eecd243c9c6b2e0156b3844c4da8aeb8e9/opentelemetry_instrumentation_openai_agents-0.52.5-py3-none-any.whl", hash = "sha256:60a489133e6b8b8d6e549c7f7ea9291a46acd0d879711ffde82a35ab11cda9dd", size = 16222, upload-time = "2026-02-23T15:40:04.662Z" }, ] +[[package]] +name = "opentelemetry-processor-baggage" +version = "0.63b1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-sdk" }, + { name = "wrapt" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/7b/43/69a8dcb2a540809c02bfc31d01b2e8228b4dd28c363d6d59577bcf9f1361/opentelemetry_processor_baggage-0.63b1.tar.gz", hash = "sha256:334b77963ea5807efd6f05664a6064aa92fc6c03571edbf1f749b9dee370d567", size = 8834, upload-time = "2026-05-21T16:36:54.591Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/23/44/3179fc05c69ca82657565d3223feb14d5208ede33e6838e42e7fe2dc01b8/opentelemetry_processor_baggage-0.63b1-py3-none-any.whl", hash = "sha256:b205c343720ce4d5e420204e09862a043917ee433b2304d87bb6f388084f3c15", size = 9488, upload-time = "2026-05-21T16:36:06.756Z" }, +] + [[package]] name = "opentelemetry-proto" version = "1.42.1" From 459863c13e21c0c8339928bc005ecea297837d51 Mon Sep 17 00:00:00 2001 From: rtemperini <29879569+rtemperini@users.noreply.github.com> Date: Mon, 7 Sep 2026 10:40:50 +0200 Subject: [PATCH 10/13] fix(tracing): reject leftover hash mappings and forward OTEL value limit Leftover {hash:...} allowlist entries must fail closed so a raw email cannot land on user.hash. Forward OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT into agent runtimes so the documented truncation control actually applies. Signed-off-by: rtemperini <29879569+rtemperini@users.noreply.github.com> Co-authored-by: Cursor --- go/adk/pkg/telemetry/context_attributes.go | 10 ++++++-- .../pkg/telemetry/context_attributes_test.go | 24 +++++++++++++++++++ go/core/internal/translator/compiler_test.go | 16 ++++++++----- go/core/internal/translator/otel.go | 1 + go/core/internal/translator/otel_test.go | 2 ++ .../core/tracing/_context_attributes.py | 13 ++++++---- .../tests/test_caller_context_attributes.py | 22 +++++++++++++++++ 7 files changed, 76 insertions(+), 12 deletions(-) diff --git a/go/adk/pkg/telemetry/context_attributes.go b/go/adk/pkg/telemetry/context_attributes.go index de44622259..b12f2424d1 100644 --- a/go/adk/pkg/telemetry/context_attributes.go +++ b/go/adk/pkg/telemetry/context_attributes.go @@ -164,12 +164,18 @@ func parseJSONAllowlist(raw string) []contextMapping { continue } var spec struct { - From string `json:"from"` - To string `json:"to"` + From string `json:"from"` + To string `json:"to"` + Hash json.RawMessage `json:"hash"` } if err := json.Unmarshal(item, &spec); err != nil { continue } + // Leftover {hash:...} config is no longer supported. Drop the + // mapping rather than stamping the raw source value onto `to`. + if spec.Hash != nil { + continue + } if mapping, ok := newContextMapping(spec.From, spec.To); ok { mappings = append(mappings, mapping) } diff --git a/go/adk/pkg/telemetry/context_attributes_test.go b/go/adk/pkg/telemetry/context_attributes_test.go index dcfc3a016d..e066028d85 100644 --- a/go/adk/pkg/telemetry/context_attributes_test.go +++ b/go/adk/pkg/telemetry/context_attributes_test.go @@ -121,6 +121,30 @@ func TestContextWithPromotedMetadata(t *testing.T) { metadata: map[string]any{"bad key": "x", "user.id": "ok"}, want: map[string]string{"user.id": "ok"}, }, + { + name: "rejects leftover hmac-sha256 hash mappings", + allowlist: `[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]`, + metadata: map[string]any{"email": "person@example.test"}, + want: map[string]string{}, + }, + { + name: "does not remap leftover hashed mappings from baggage", + allowlist: `[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]`, + baggageVals: map[string]string{"email": "person@example.test"}, + want: map[string]string{"email": "person@example.test"}, + }, + { + name: "rejects unknown leftover hash mappings", + allowlist: `[{"from":"email","to":"user.hash","hash":"sha3"}]`, + metadata: map[string]any{"email": "person@example.test"}, + want: map[string]string{}, + }, + { + name: "keeps mappings without hash when a hashed entry is dropped", + allowlist: `[{"from":"email","to":"user.hash","hash":"hmac-sha256"},{"from":"sub","to":"user.id"}]`, + metadata: map[string]any{"email": "person@example.test", "sub": "opaque-subject"}, + want: map[string]string{"user.id": "opaque-subject"}, + }, } for _, tt := range tests { diff --git a/go/core/internal/translator/compiler_test.go b/go/core/internal/translator/compiler_test.go index f7bc1392d9..e7f96c8220 100644 --- a/go/core/internal/translator/compiler_test.go +++ b/go/core/internal/translator/compiler_test.go @@ -295,6 +295,7 @@ func TestCompileAgentTemplateResolvesCredentialsForSubstrate(t *testing.T) { func TestCompileAgentTemplateForwardsOtelEnvironment(t *testing.T) { t.Setenv("OTEL_TRACING_ENABLED", "true") + t.Setenv("OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT", "256") harness := &v1alpha3.Harness{ ObjectMeta: metav1.ObjectMeta{Name: "kagent", Namespace: "test"}, Spec: v1alpha3.HarnessSpec{ @@ -318,15 +319,18 @@ func TestCompileAgentTemplateForwardsOtelEnvironment(t *testing.T) { if err != nil { t.Fatal(err) } + got := map[string]string{} for _, variable := range spec.Environment { - if variable.Name == "OTEL_TRACING_ENABLED" { - if variable.Value != "true" { - t.Fatalf("OTEL_TRACING_ENABLED = %q, want %q", variable.Value, "true") - } - return + if variable.Name == "OTEL_TRACING_ENABLED" || variable.Name == "OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT" { + got[variable.Name] = variable.Value } } - t.Fatalf("OTEL_TRACING_ENABLED missing from runtime revision environment: %+v", spec.Environment) + if got["OTEL_TRACING_ENABLED"] != "true" { + t.Fatalf("OTEL_TRACING_ENABLED = %q, want %q (env: %+v)", got["OTEL_TRACING_ENABLED"], "true", spec.Environment) + } + if got["OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT"] != "256" { + t.Fatalf("OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT = %q, want %q (env: %+v)", got["OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT"], "256", spec.Environment) + } } func TestCompileAgentTemplateSharedAgent(t *testing.T) { diff --git a/go/core/internal/translator/otel.go b/go/core/internal/translator/otel.go index 60dcbdc6bb..a505eb608a 100644 --- a/go/core/internal/translator/otel.go +++ b/go/core/internal/translator/otel.go @@ -16,6 +16,7 @@ var otelEnvNames = []string{ "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", "OTEL_EXPORTER_OTLP_PROTOCOL", "OTEL_EXPORTER_OTLP_TRACES_PROTOCOL", + "OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT", } // OtelEnvFromProcess returns the controller's supported tracing configuration diff --git a/go/core/internal/translator/otel_test.go b/go/core/internal/translator/otel_test.go index ea6a4ec0eb..fe0381acf2 100644 --- a/go/core/internal/translator/otel_test.go +++ b/go/core/internal/translator/otel_test.go @@ -14,6 +14,7 @@ func TestOtelEnvFromProcess(t *testing.T) { t.Setenv("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", "http://collector:4317") t.Setenv("OTEL_EXPORTER_OTLP_PROTOCOL", "http/protobuf") t.Setenv("OTEL_EXPORTER_OTLP_TRACES_PROTOCOL", "grpc") + t.Setenv("OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT", "256") t.Setenv("OTEL_EXPORTER_OTLP_HEADERS", "authorization=secret") t.Setenv("OTEL_RESOURCE_ATTRIBUTES", "service.name=controller") t.Setenv("OTEL_SERVICE_NAME", "controller") @@ -25,6 +26,7 @@ func TestOtelEnvFromProcess(t *testing.T) { {Name: "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", Value: "http://collector:4317"}, {Name: "OTEL_EXPORTER_OTLP_PROTOCOL", Value: "http/protobuf"}, {Name: "OTEL_EXPORTER_OTLP_TRACES_PROTOCOL", Value: "grpc"}, + {Name: "OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT", Value: "256"}, } if !reflect.DeepEqual(got, want) { t.Errorf("OtelEnvFromProcess() = %#v, want %#v", got, want) diff --git a/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py b/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py index 391a31a2c3..df9e521422 100644 --- a/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py +++ b/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py @@ -142,12 +142,17 @@ def _parse_json_allowlist(raw: str) -> list[_ContextMapping]: if isinstance(item, str): mapping = _new_context_mapping(item, "") elif isinstance(item, dict): - from_key, from_ok = _json_string_field(item, "from") - to_key, to_ok = _json_string_field(item, "to") - if not from_ok or not to_ok: + # Leftover {hash:...} config is no longer supported. Drop the + # mapping rather than stamping the raw source value onto `to`. + if "hash" in item: mapping = None else: - mapping = _new_context_mapping(from_key, to_key) + from_key, from_ok = _json_string_field(item, "from") + to_key, to_ok = _json_string_field(item, "to") + if not from_ok or not to_ok: + mapping = None + else: + mapping = _new_context_mapping(from_key, to_key) else: mapping = None if mapping is not None: diff --git a/python/packages/kagent-core/tests/test_caller_context_attributes.py b/python/packages/kagent-core/tests/test_caller_context_attributes.py index e2637c35cf..d0df4d7c60 100644 --- a/python/packages/kagent-core/tests/test_caller_context_attributes.py +++ b/python/packages/kagent-core/tests/test_caller_context_attributes.py @@ -127,6 +127,28 @@ def test_drops_keys_with_whitespace(self, monkeypatch): set_allowlist(monkeypatch, "bad key,user.id") assert baggage_values(context_with_promoted_metadata({"bad key": "x", "user.id": "ok"})) == {"user.id": "ok"} + def test_rejects_leftover_hmac_sha256_hash_mappings(self, monkeypatch): + set_allowlist(monkeypatch, '[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]') + assert baggage_values(context_with_promoted_metadata({"email": "person@example.test"})) == {} + + def test_does_not_remap_leftover_hashed_mappings_from_baggage(self, monkeypatch): + set_allowlist(monkeypatch, '[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]') + ctx = baggage_context({"email": "person@example.test"}) + assert baggage_values(context_with_promoted_metadata({}, ctx)) == {"email": "person@example.test"} + + def test_rejects_unknown_leftover_hash_mappings(self, monkeypatch): + set_allowlist(monkeypatch, '[{"from":"email","to":"user.hash","hash":"sha3"}]') + assert baggage_values(context_with_promoted_metadata({"email": "person@example.test"})) == {} + + def test_keeps_mappings_without_hash_when_hashed_entry_dropped(self, monkeypatch): + set_allowlist( + monkeypatch, + '[{"from":"email","to":"user.hash","hash":"hmac-sha256"},{"from":"sub","to":"user.id"}]', + ) + assert baggage_values( + context_with_promoted_metadata({"email": "person@example.test", "sub": "opaque-subject"}) + ) == {"user.id": "opaque-subject"} + def test_caps_the_allowlist(self, monkeypatch): keys = [f"k{i}" for i in range(MAX_CONTEXT_KEYS + 4)] set_allowlist(monkeypatch, ",".join(keys)) From 074c5f2d5f264ebdf72235f02c3126fadfefdf1a Mon Sep 17 00:00:00 2001 From: rtemperini <29879569+rtemperini@users.noreply.github.com> Date: Mon, 21 Sep 2026 15:29:11 +0200 Subject: [PATCH 11/13] fix(tracing): suppress rejected mapping sources on metadata spans Leftover {hash:...} allowlist entries were dropped before their from key entered the suppression set, so the executor still stamped the raw value as a2a.message.metadata.. Keep those sources covered without re-implementing hashing. Signed-off-by: rtemperini <29879569+rtemperini@users.noreply.github.com> --- go/adk/pkg/telemetry/attributes_test.go | 38 ++++++++++++++++++++++ go/adk/pkg/telemetry/context_attributes.go | 29 +++++++++++++---- 2 files changed, 61 insertions(+), 6 deletions(-) diff --git a/go/adk/pkg/telemetry/attributes_test.go b/go/adk/pkg/telemetry/attributes_test.go index 46e790d068..d5b3c95d3a 100644 --- a/go/adk/pkg/telemetry/attributes_test.go +++ b/go/adk/pkg/telemetry/attributes_test.go @@ -234,6 +234,44 @@ func TestSetMessageMetadataAttributes_SkipsAllowlistedSources(t *testing.T) { } } +// The executor stamps leftover A2A metadata via SetMessageMetadataAttributes. +// A rejected leftover {hash:...} mapping must still suppress its source so +// the raw value cannot appear as a2a.message.metadata.. +func TestSetMessageMetadataAttributes_SkipsRejectedHashSources(t *testing.T) { + setAllowlist(t, `[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]`) + + attrs := recordMessageMetadata(t, map[string]any{ + "email": "person@example.test", + "channel": "C0AB1", + }) + + if _, exists := attrs["a2a.message.metadata.email"]; exists { + t.Fatal("rejected hash mapping source must not be stamped as a2a.message.metadata.email") + } + if _, exists := attrs["user.hash"]; exists { + t.Fatal("rejected hash mappings must not emit user.hash") + } + if got := attrs["a2a.message.metadata.channel"].AsString(); got != "C0AB1" { + t.Errorf("non-allowlisted channel = %q, want C0AB1", got) + } +} + +func TestSetMessageMetadataAttributes_SkipsInvalidMappingSources(t *testing.T) { + setAllowlist(t, `[{"from":"email","to":"bad key"}]`) + + attrs := recordMessageMetadata(t, map[string]any{ + "email": "person@example.test", + "channel": "C0AB1", + }) + + if _, exists := attrs["a2a.message.metadata.email"]; exists { + t.Fatal("invalid mapping source must not be stamped as a2a.message.metadata.email") + } + if got := attrs["a2a.message.metadata.channel"].AsString(); got != "C0AB1" { + t.Errorf("non-allowlisted channel = %q, want C0AB1", got) + } +} + func recordMessageMetadata(t *testing.T, metadata map[string]any) map[string]attribute.Value { t.Helper() exporter := tracetest.NewInMemoryExporter() diff --git a/go/adk/pkg/telemetry/context_attributes.go b/go/adk/pkg/telemetry/context_attributes.go index b12f2424d1..d9a524ab11 100644 --- a/go/adk/pkg/telemetry/context_attributes.go +++ b/go/adk/pkg/telemetry/context_attributes.go @@ -33,7 +33,9 @@ type contextMapping struct { // contextPolicy is the parsed allowlist plus every source key named in it. // Source keys are skipped by SetMessageMetadataAttributes so allowlisted -// values are not also stamped as a2a.message.metadata.. +// values are not also stamped as a2a.message.metadata.. Rejected +// leftover {hash:...} mappings still contribute their from key so the +// raw source cannot leak on that generic metadata path. type contextPolicy struct { mappings []contextMapping sources map[string]struct{} @@ -57,8 +59,8 @@ func policyCoveredContextSources() map[string]struct{} { } func parseContextPolicy() contextPolicy { - mappings := parseAllowedContextMappings() - sources := make(map[string]struct{}, len(mappings)) + sources := map[string]struct{}{} + mappings := parseAllowedContextMappings(sources) for _, mapping := range mappings { sources[mapping.source] = struct{}{} } @@ -128,13 +130,13 @@ func AllowedBaggageCopyFilter() baggagecopy.Filter { } } -func parseAllowedContextMappings() []contextMapping { +func parseAllowedContextMappings(sources map[string]struct{}) []contextMapping { raw := strings.TrimSpace(os.Getenv(traceContextKeysEnvVar)) if raw == "" { return nil } if strings.HasPrefix(raw, "[") { - return capMappings(parseJSONAllowlist(raw)) + return capMappings(parseJSONAllowlist(raw, sources)) } return capMappings(parseCommaAllowlist(raw)) } @@ -149,7 +151,7 @@ func parseCommaAllowlist(raw string) []contextMapping { return mappings } -func parseJSONAllowlist(raw string) []contextMapping { +func parseJSONAllowlist(raw string, sources map[string]struct{}) []contextMapping { var items []json.RawMessage if err := json.Unmarshal([]byte(raw), &items); err != nil { return nil @@ -173,16 +175,31 @@ func parseJSONAllowlist(raw string) []contextMapping { } // Leftover {hash:...} config is no longer supported. Drop the // mapping rather than stamping the raw source value onto `to`. + // Keep `from` in the metadata-suppression set so the executor + // cannot write a2a.message.metadata. either. if spec.Hash != nil { + coverContextSource(sources, spec.From) continue } if mapping, ok := newContextMapping(spec.From, spec.To); ok { mappings = append(mappings, mapping) + continue } + coverContextSource(sources, spec.From) } return mappings } +// coverContextSource records a named from key so SetMessageMetadataAttributes +// skips it, even when the mapping itself was rejected. +func coverContextSource(sources map[string]struct{}, from string) { + from = strings.TrimSpace(from) + if from == "" || len([]rune(from)) > maxContextKeyLength || !isAttributeKey(from) { + return + } + sources[from] = struct{}{} +} + func newContextMapping(from, to string) (contextMapping, bool) { from = strings.TrimSpace(from) to = strings.TrimSpace(to) From a3e8c4fb716b0cf8192e8f14f7364a63a92c2656 Mon Sep 17 00:00:00 2001 From: rtemperini <29879569+rtemperini@users.noreply.github.com> Date: Mon, 19 Oct 2026 09:49:01 +0200 Subject: [PATCH 12/13] fix(tracing): enforce allowlisted metadata export Signed-off-by: rtemperini <29879569+rtemperini@users.noreply.github.com> --- docs/architecture/trace-context.md | 4 +- go/adk/pkg/a2a/executor.go | 1 - go/adk/pkg/a2a/executor_test.go | 39 +++++ go/adk/pkg/telemetry/attributes.go | 45 ------ go/adk/pkg/telemetry/attributes_test.go | 147 ------------------ go/adk/pkg/telemetry/context_attributes.go | 54 +------ .../translator/kagent/context_policy_test.go | 41 +++++ .../tests/controller-deployment_test.yaml | 6 +- .../core/tracing/_context_attributes.py | 4 +- .../tests/test_caller_context_attributes.py | 16 +- .../src/kagent/openai/_agent_executor.py | 1 + .../tests/test_caller_context.py | 41 ++++- 12 files changed, 149 insertions(+), 250 deletions(-) create mode 100644 go/core/internal/translator/kagent/context_policy_test.go diff --git a/docs/architecture/trace-context.md b/docs/architecture/trace-context.md index 6a3d68f2ea..f99e40f83b 100644 --- a/docs/architecture/trace-context.md +++ b/docs/architecture/trace-context.md @@ -15,12 +15,12 @@ allowlist. | Setting | Default | Description | |---|---|---| -| Helm `otel.tracing.contextKeys` | `[]` | List of baggage keys, or `{from, to}` mappings, to copy onto spans | +| Helm `otel.traces.contextKeys` | `[]` | List of baggage keys, or `{from, to}` mappings, to copy onto spans | | Env `KAGENT_TRACE_CONTEXT_KEYS` | `""` | Comma-separated keys, or a JSON array of the same mappings | ```yaml otel: - tracing: + traces: enabled: true contextKeys: - {from: sub, to: user.id} diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index 901cc4eaf1..49cbe5cc42 100644 --- a/go/adk/pkg/a2a/executor.go +++ b/go/adk/pkg/a2a/executor.go @@ -259,7 +259,6 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.ExecutorCon tracing.InvocationFromContext(ctx).SetAttributes(tracing.RequestIdentity(sessionID, string(reqCtx.TaskID), resumed)...) ctx = telemetry.ContextWithPromotedMetadata(ctx, reqCtx.Message.Metadata) ctx = telemetry.WithRequestAttributes(ctx, requestSpanAttributes(sessionID, string(reqCtx.TaskID), trustedUserID)...) - telemetry.SetMessageMetadataAttributes(ctx, reqCtx.Message.Metadata) e.logger.InfoContext(ctx, "execute", "task_id", reqCtx.TaskID, diff --git a/go/adk/pkg/a2a/executor_test.go b/go/adk/pkg/a2a/executor_test.go index 37e80f492f..21f6832c0f 100644 --- a/go/adk/pkg/a2a/executor_test.go +++ b/go/adk/pkg/a2a/executor_test.go @@ -910,3 +910,42 @@ func TestRequestSpanAttributesCarryOnlyATrustedUser(t *testing.T) { }) } } + +func TestKAgentExecutorDoesNotExportUnallowlistedMetadata(t *testing.T) { + t.Setenv("KAGENT_TRACE_CONTEXT_KEYS", "") + exporter, flush := installRecordingTracer(t) + ctx, root := otel.Tracer("test").Start(t.Context(), "invocation") + message := a2atype.NewMessage(a2atype.MessageRoleUser, a2atype.NewTextPart("hello")) + message.Metadata = map[string]any{ + "private_token": "must-not-be-exported", + "attempt": float64(3), + "enabled": true, + } + executor := &KAgentExecutor{ + builtin: &recordingExecutor{span: "generate_content"}, + logger: slog.New(slog.DiscardHandler), + } + for _, err := range executor.Execute(ctx, &a2asrv.ExecutorContext{ + ContextID: "ctx-1", TaskID: "task-1", Message: message, + }) { + if err != nil { + t.Fatal(err) + } + } + root.End() + if err := flush(t.Context()); err != nil { + t.Fatal(err) + } + spans := exporter.GetSpans() + if len(spans) != 2 { + t.Fatalf("exported %d spans, want invocation and model spans", len(spans)) + } + for _, span := range spans { + for _, attr := range span.Attributes { + key := string(attr.Key) + if strings.HasPrefix(key, "a2a.message.metadata.") || key == "private_token" || key == "attempt" || key == "enabled" { + t.Errorf("%s exported unallowlisted metadata attribute %q", span.Name, key) + } + } + } +} diff --git a/go/adk/pkg/telemetry/attributes.go b/go/adk/pkg/telemetry/attributes.go index eebb75fdaa..b2f870106f 100644 --- a/go/adk/pkg/telemetry/attributes.go +++ b/go/adk/pkg/telemetry/attributes.go @@ -2,11 +2,9 @@ package telemetry import ( "context" - "fmt" "go.opentelemetry.io/otel/attribute" sdktrace "go.opentelemetry.io/otel/sdk/trace" - "go.opentelemetry.io/otel/trace" ) type requestAttributesKey struct{} @@ -33,46 +31,3 @@ func (requestAttributesSpanProcessor) OnEnd(sdktrace.ReadOnlySpan) {} func (requestAttributesSpanProcessor) Shutdown(context.Context) error { return nil } func (requestAttributesSpanProcessor) ForceFlush(context.Context) error { return nil } - -// SetMessageMetadataAttributes sets scalar values from an A2A message's -// metadata as span attributes on the current span. -// -// Keys named in KAGENT_TRACE_CONTEXT_KEYS are skipped. Those values are -// promoted into baggage and copied onto spans by baggagecopy; stamping -// them here would also emit a2a.message.metadata.. -func SetMessageMetadataAttributes(ctx context.Context, metadata map[string]any) { - if len(metadata) == 0 { - return - } - covered := policyCoveredContextSources() - var attrs []attribute.KeyValue - for k, v := range metadata { - if _, skip := covered[k]; skip { - continue - } - key := "a2a.message.metadata." + k - switch val := v.(type) { - case string: - if val != "" { - attrs = append(attrs, attribute.String(key, val)) - } - case bool: - attrs = append(attrs, attribute.Bool(key, val)) - case float64: - attrs = append(attrs, attribute.String(key, fmt.Sprintf("%g", val))) - case int: - attrs = append(attrs, attribute.Int(key, val)) - case int64: - attrs = append(attrs, attribute.Int64(key, val)) - } - } - setSpanAttributes(ctx, attrs...) -} - -func setSpanAttributes(ctx context.Context, attrs ...attribute.KeyValue) { - span := trace.SpanFromContext(ctx) - if !span.IsRecording() || len(attrs) == 0 { - return - } - span.SetAttributes(attrs...) -} diff --git a/go/adk/pkg/telemetry/attributes_test.go b/go/adk/pkg/telemetry/attributes_test.go index a0156658e6..0ec31e6f1a 100644 --- a/go/adk/pkg/telemetry/attributes_test.go +++ b/go/adk/pkg/telemetry/attributes_test.go @@ -2,7 +2,6 @@ package telemetry import ( "context" - "strings" "testing" "go.opentelemetry.io/otel/attribute" @@ -49,152 +48,6 @@ func TestWithRequestAttributesKeepsContextWhenEmpty(t *testing.T) { t.Fatal("no attributes should return the same context") } } -func TestSetMessageMetadataAttributes(t *testing.T) { - exporter := tracetest.NewInMemoryExporter() - tp := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter)) - t.Cleanup(func() { _ = tp.Shutdown(context.Background()) }) - - tracer := tp.Tracer("test") - ctx, span := tracer.Start(context.Background(), "test-span") - - SetMessageMetadataAttributes(ctx, map[string]any{ - "approver_email": "admin@example.com", - "attempt_count": float64(3), - "dry_run": true, - "nested": map[string]any{"should": "be skipped"}, - "list_val": []string{"also", "skipped"}, - "empty_str": "", - }) - span.End() - - spans := exporter.GetSpans() - if len(spans) == 0 { - t.Fatal("no spans recorded") - } - attrs := make(map[string]attribute.Value) - for _, a := range spans[0].Attributes { - attrs[string(a.Key)] = a.Value - } - - if got := attrs["a2a.message.metadata.approver_email"].AsString(); got != "admin@example.com" { - t.Errorf("approver_email: got %q, want %q", got, "admin@example.com") - } - if got := attrs["a2a.message.metadata.attempt_count"].AsString(); got != "3" { - t.Errorf("attempt_count: got %q, want %q", got, "3") - } - if got := attrs["a2a.message.metadata.dry_run"].AsBool(); !got { - t.Errorf("dry_run: got %v, want true", got) - } - if _, exists := attrs["a2a.message.metadata.nested"]; exists { - t.Error("nested map should not be set as span attribute") - } - if _, exists := attrs["a2a.message.metadata.list_val"]; exists { - t.Error("list value should not be set as span attribute") - } - if _, exists := attrs["a2a.message.metadata.empty_str"]; exists { - t.Error("empty string should not be set as span attribute") - } -} - -func TestSetMessageMetadataAttributes_SkipsAllowlistedSources(t *testing.T) { - setAllowlist(t, `[{"from":"sub","to":"user.id"}]`) - - attrs := recordMessageMetadata(t, map[string]any{ - "sub": "opaque-subject", - "channel": "C0AB1", - }) - - if _, exists := attrs["a2a.message.metadata.sub"]; exists { - t.Fatal("allowlisted source must not be stamped as a2a.message.metadata.sub") - } - if got := attrs["a2a.message.metadata.channel"].AsString(); got != "C0AB1" { - t.Errorf("non-allowlisted channel = %q, want C0AB1", got) - } -} - -// The executor stamps leftover A2A metadata via SetMessageMetadataAttributes. -// A rejected leftover {hash:...} mapping must still suppress its source so -// the raw value cannot appear as a2a.message.metadata.. -func TestSetMessageMetadataAttributes_SkipsRejectedHashSources(t *testing.T) { - setAllowlist(t, `[{"from":"email","to":"user.hash","hash":"hmac-sha256"}]`) - - attrs := recordMessageMetadata(t, map[string]any{ - "email": "person@example.test", - "channel": "C0AB1", - }) - - if _, exists := attrs["a2a.message.metadata.email"]; exists { - t.Fatal("rejected hash mapping source must not be stamped as a2a.message.metadata.email") - } - if _, exists := attrs["user.hash"]; exists { - t.Fatal("rejected hash mappings must not emit user.hash") - } - if got := attrs["a2a.message.metadata.channel"].AsString(); got != "C0AB1" { - t.Errorf("non-allowlisted channel = %q, want C0AB1", got) - } -} - -func TestSetMessageMetadataAttributes_SkipsInvalidMappingSources(t *testing.T) { - setAllowlist(t, `[{"from":"email","to":"bad key"}]`) - - attrs := recordMessageMetadata(t, map[string]any{ - "email": "person@example.test", - "channel": "C0AB1", - }) - - if _, exists := attrs["a2a.message.metadata.email"]; exists { - t.Fatal("invalid mapping source must not be stamped as a2a.message.metadata.email") - } - if got := attrs["a2a.message.metadata.channel"].AsString(); got != "C0AB1" { - t.Errorf("non-allowlisted channel = %q, want C0AB1", got) - } -} - -func recordMessageMetadata(t *testing.T, metadata map[string]any) map[string]attribute.Value { - t.Helper() - exporter := tracetest.NewInMemoryExporter() - tp := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter)) - t.Cleanup(func() { _ = tp.Shutdown(context.Background()) }) - - tracer := tp.Tracer("test") - ctx, span := tracer.Start(context.Background(), "test-span") - SetMessageMetadataAttributes(ctx, metadata) - span.End() - - spans := exporter.GetSpans() - if len(spans) == 0 { - t.Fatal("no spans recorded") - } - attrs := make(map[string]attribute.Value) - for _, a := range spans[0].Attributes { - attrs[string(a.Key)] = a.Value - } - return attrs -} - -func TestSetMessageMetadataAttributes_NilAndEmpty(t *testing.T) { - exporter := tracetest.NewInMemoryExporter() - tp := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter)) - t.Cleanup(func() { _ = tp.Shutdown(context.Background()) }) - - tracer := tp.Tracer("test") - ctx, span := tracer.Start(context.Background(), "test-span") - - SetMessageMetadataAttributes(ctx, nil) - SetMessageMetadataAttributes(ctx, map[string]any{}) - span.End() - - spans := exporter.GetSpans() - if len(spans) == 0 { - t.Fatal("no spans recorded") - } - for _, a := range spans[0].Attributes { - if strings.HasPrefix(string(a.Key), "a2a.message.metadata.") { - t.Errorf("no metadata attributes expected, got %q", a.Key) - } - } -} - func spanAttributesByName(t *testing.T, spans tracetest.SpanStubs, name string) map[string]attribute.Value { t.Helper() diff --git a/go/adk/pkg/telemetry/context_attributes.go b/go/adk/pkg/telemetry/context_attributes.go index d9a524ab11..6bd8923296 100644 --- a/go/adk/pkg/telemetry/context_attributes.go +++ b/go/adk/pkg/telemetry/context_attributes.go @@ -31,40 +31,16 @@ type contextMapping struct { attribute string } -// contextPolicy is the parsed allowlist plus every source key named in it. -// Source keys are skipped by SetMessageMetadataAttributes so allowlisted -// values are not also stamped as a2a.message.metadata.. Rejected -// leftover {hash:...} mappings still contribute their from key so the -// raw source cannot leak on that generic metadata path. -type contextPolicy struct { - mappings []contextMapping - sources map[string]struct{} -} - -// loadContextPolicy parses KAGENT_TRACE_CONTEXT_KEYS once. The allowlist -// cannot change without a process restart; tests call -// resetAllowedContextMappings after t.Setenv. -var loadContextPolicy = sync.OnceValue(parseContextPolicy) +// loadContextMappings parses the allowlist once per process. Tests reset it +// after changing the environment. +var loadContextMappings = sync.OnceValue(parseAllowedContextMappings) func resetAllowedContextMappings() { - loadContextPolicy = sync.OnceValue(parseContextPolicy) + loadContextMappings = sync.OnceValue(parseAllowedContextMappings) } func allowedContextMappings() []contextMapping { - return loadContextPolicy().mappings -} - -func policyCoveredContextSources() map[string]struct{} { - return loadContextPolicy().sources -} - -func parseContextPolicy() contextPolicy { - sources := map[string]struct{}{} - mappings := parseAllowedContextMappings(sources) - for _, mapping := range mappings { - sources[mapping.source] = struct{}{} - } - return contextPolicy{mappings: mappings, sources: sources} + return loadContextMappings() } // ContextWithPromotedMetadata copies allowlisted A2A message.metadata values @@ -130,13 +106,13 @@ func AllowedBaggageCopyFilter() baggagecopy.Filter { } } -func parseAllowedContextMappings(sources map[string]struct{}) []contextMapping { +func parseAllowedContextMappings() []contextMapping { raw := strings.TrimSpace(os.Getenv(traceContextKeysEnvVar)) if raw == "" { return nil } if strings.HasPrefix(raw, "[") { - return capMappings(parseJSONAllowlist(raw, sources)) + return capMappings(parseJSONAllowlist(raw)) } return capMappings(parseCommaAllowlist(raw)) } @@ -151,7 +127,7 @@ func parseCommaAllowlist(raw string) []contextMapping { return mappings } -func parseJSONAllowlist(raw string, sources map[string]struct{}) []contextMapping { +func parseJSONAllowlist(raw string) []contextMapping { var items []json.RawMessage if err := json.Unmarshal([]byte(raw), &items); err != nil { return nil @@ -175,31 +151,17 @@ func parseJSONAllowlist(raw string, sources map[string]struct{}) []contextMappin } // Leftover {hash:...} config is no longer supported. Drop the // mapping rather than stamping the raw source value onto `to`. - // Keep `from` in the metadata-suppression set so the executor - // cannot write a2a.message.metadata. either. if spec.Hash != nil { - coverContextSource(sources, spec.From) continue } if mapping, ok := newContextMapping(spec.From, spec.To); ok { mappings = append(mappings, mapping) continue } - coverContextSource(sources, spec.From) } return mappings } -// coverContextSource records a named from key so SetMessageMetadataAttributes -// skips it, even when the mapping itself was rejected. -func coverContextSource(sources map[string]struct{}, from string) { - from = strings.TrimSpace(from) - if from == "" || len([]rune(from)) > maxContextKeyLength || !isAttributeKey(from) { - return - } - sources[from] = struct{}{} -} - func newContextMapping(from, to string) (contextMapping, bool) { from = strings.TrimSpace(from) to = strings.TrimSpace(to) diff --git a/go/core/internal/translator/kagent/context_policy_test.go b/go/core/internal/translator/kagent/context_policy_test.go new file mode 100644 index 0000000000..aafc479324 --- /dev/null +++ b/go/core/internal/translator/kagent/context_policy_test.go @@ -0,0 +1,41 @@ +package kagent + +import ( + "testing" + + "github.com/kagent-dev/kagent/go/core/pkg/env" + corev1 "k8s.io/api/core/v1" +) + +func TestApplyOperatorOnlyTraceContextPolicy(t *testing.T) { + name := env.KagentTraceContextKeys.Name() + for _, operatorValue := range []string{"", "user.id"} { + t.Run("operator="+operatorValue, func(t *testing.T) { + t.Setenv(name, operatorValue) + values := applyOperatorOnlyEnv([]corev1.EnvVar{ + {Name: "OTHER", Value: "kept"}, + {Name: name, Value: "private_token"}, + {Name: name, ValueFrom: &corev1.EnvVarSource{ + SecretKeyRef: &corev1.SecretKeySelector{Key: "tenant-policy"}, + }}, + }, env.KagentTraceContextKeys) + count := 0 + for _, value := range values { + if value.Name != name { + continue + } + count++ + if value.Value != operatorValue || value.ValueFrom != nil { + t.Fatalf("tenant policy survived: %#v", value) + } + } + want := 0 + if operatorValue != "" { + want = 1 + } + if count != want || values[0].Name != "OTHER" || values[0].Value != "kept" { + t.Fatalf("environment = %#v, want %d operator policy entries and OTHER kept", values, want) + } + }) + } +} diff --git a/helm/kagent/tests/controller-deployment_test.yaml b/helm/kagent/tests/controller-deployment_test.yaml index 4fe90bbf64..13b712e246 100644 --- a/helm/kagent/tests/controller-deployment_test.yaml +++ b/helm/kagent/tests/controller-deployment_test.yaml @@ -986,10 +986,10 @@ tests: - notExists: path: data.KAGENT_TRACE_CONTEXT_KEYS - - it: should join otel.tracing.contextKeys into KAGENT_TRACE_CONTEXT_KEYS + - it: should join otel.traces.contextKeys into KAGENT_TRACE_CONTEXT_KEYS template: controller-configmap.yaml set: - otel.tracing.contextKeys: + otel.traces.contextKeys: - sub - thread_id - channel @@ -1001,7 +1001,7 @@ tests: - it: should encode contextKeys mappings as JSON template: controller-configmap.yaml set: - otel.tracing.contextKeys: + otel.traces.contextKeys: - from: sub to: user.id - from: thread_id diff --git a/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py b/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py index df9e521422..aafc380b5c 100644 --- a/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py +++ b/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py @@ -13,7 +13,7 @@ from opentelemetry import baggage from opentelemetry import context as otel_context -from kagent.core.a2a._consts import read_message_metadata +from google.protobuf.json_format import MessageToDict # Allowlist of caller-supplied context keys copied from W3C baggage onto # agent spans. Unset or empty (the default) disables promotion entirely. @@ -50,7 +50,7 @@ def context_with_promoted_metadata( if not mappings: return ctx if metadata is None and message is not None: - metadata = read_message_metadata(message) + metadata = MessageToDict(message.metadata) if message.HasField("metadata") else {} if metadata is None: metadata = {} diff --git a/python/packages/kagent-core/tests/test_caller_context_attributes.py b/python/packages/kagent-core/tests/test_caller_context_attributes.py index d0df4d7c60..53f4e0facb 100644 --- a/python/packages/kagent-core/tests/test_caller_context_attributes.py +++ b/python/packages/kagent-core/tests/test_caller_context_attributes.py @@ -1,6 +1,8 @@ from unittest.mock import MagicMock import pytest +from a2a.types import Message, Role +from google.protobuf.json_format import ParseDict from opentelemetry import baggage from opentelemetry import context as otel_context from opentelemetry.processor.baggage import BaggageSpanProcessor @@ -158,7 +160,7 @@ def test_caps_the_allowlist(self, monkeypatch): def test_skips_message_decode_when_allowlist_empty(self, monkeypatch): set_allowlist(monkeypatch, "") decode = MagicMock(return_value={"thread_id": "T1"}) - monkeypatch.setattr("kagent.core.tracing._context_attributes.read_message_metadata", decode) + monkeypatch.setattr("kagent.core.tracing._context_attributes.MessageToDict", decode) assert baggage_values(context_with_promoted_metadata(message=MagicMock())) == {} decode.assert_not_called() @@ -166,7 +168,7 @@ def test_decodes_message_after_allowlist_check(self, monkeypatch): set_allowlist(monkeypatch, "thread_id") message = MagicMock() monkeypatch.setattr( - "kagent.core.tracing._context_attributes.read_message_metadata", + "kagent.core.tracing._context_attributes.MessageToDict", lambda _message: {"thread_id": "T1"}, ) assert baggage_values(context_with_promoted_metadata(message=message)) == {"thread_id": "T1"} @@ -251,3 +253,13 @@ def test_runtime_attributes_win_over_allowlisted_baggage(monkeypatch): attrs = _span_attrs(exporter.get_finished_spans(), "generate_content") assert attrs["gen_ai.conversation.id"] == "from-runtime" + + +def test_promotes_real_protobuf_message_metadata(monkeypatch): + set_allowlist(monkeypatch, '[{"from":"sub","to":"user.id"},"attempt"]') + message = Message(role=Role.ROLE_USER, message_id="msg-1") + ParseDict({"sub": "opaque-subject", "attempt": 3, "private_token": "must-not-be-exported"}, message.metadata) + assert baggage_values(context_with_promoted_metadata(message=message)) == { + "user.id": "opaque-subject", + "attempt": "3", + } diff --git a/python/packages/kagent-openai/src/kagent/openai/_agent_executor.py b/python/packages/kagent-openai/src/kagent/openai/_agent_executor.py index ee33b5ce39..b7eab340f1 100644 --- a/python/packages/kagent-openai/src/kagent/openai/_agent_executor.py +++ b/python/packages/kagent-openai/src/kagent/openai/_agent_executor.py @@ -37,6 +37,7 @@ from agents.run import Runner from kagent.core.a2a import now_timestamp from kagent.core.tracing import detach_promoted_metadata, promote_message_metadata_to_baggage +from kagent.core.tracing._span_processor import clear_kagent_span_attributes, set_kagent_span_attributes from pydantic import BaseModel from ._event_converter import convert_openai_event_to_a2a_events diff --git a/python/packages/kagent-openai/tests/test_caller_context.py b/python/packages/kagent-openai/tests/test_caller_context.py index 5a564982ff..baf4061838 100644 --- a/python/packages/kagent-openai/tests/test_caller_context.py +++ b/python/packages/kagent-openai/tests/test_caller_context.py @@ -1,9 +1,16 @@ -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock + +import pytest +from google.protobuf.json_format import ParseDict +from opentelemetry import baggage +from opentelemetry import context as otel_context +from kagent.core.tracing._context_attributes import _allowed_context_mappings +from kagent.core.tracing._span_processor import KAGENT_ATTRIBUTES_KEY from a2a.server.agent_execution.context import RequestContext from a2a.types import Message, Part, Role, SendMessageRequest -from kagent.openai._agent_executor import _convert_a2a_request_to_span_attributes +from kagent.openai._agent_executor import OpenAIAgentExecutor, _convert_a2a_request_to_span_attributes def _request_context() -> RequestContext: @@ -23,3 +30,33 @@ def test_span_attributes_are_runtime_only(): assert attrs["gen_ai.conversation.id"] == "ctx-1" assert "kagent.context.thread_id" not in attrs assert "user.id" not in attrs + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fails", [False, True]) +async def test_execute_promotes_metadata_and_restores_context(monkeypatch, fails): + monkeypatch.setenv("KAGENT_TRACE_CONTEXT_KEYS", "thread_id") + _allowed_context_mappings.cache_clear() + context = _request_context() + ParseDict({"thread_id": "T1", "private_token": "must-not-be-exported"}, context.message.metadata) + executor = OpenAIAgentExecutor(agent=MagicMock(), app_name="test") + original_context = otel_context.get_current() + + async def run_agent(request, queue): + assert baggage.get_baggage("thread_id") == "T1" + assert baggage.get_baggage("private_token") is None + assert otel_context.get_value(KAGENT_ATTRIBUTES_KEY)["gen_ai.conversation.id"] == "ctx-1" + if fails: + raise RuntimeError("test failure") + + executor._execute_agent = AsyncMock(side_effect=run_agent) + try: + if fails: + with pytest.raises(RuntimeError, match="test failure"): + await executor.execute(context, MagicMock()) + else: + await executor.execute(context, MagicMock()) + executor._execute_agent.assert_awaited_once() + assert otel_context.get_current() is original_context + finally: + _allowed_context_mappings.cache_clear() From dda3dedc2fd3569e2788f310e75d97a03d161894 Mon Sep 17 00:00:00 2001 From: rtemperini <29879569+rtemperini@users.noreply.github.com> Date: Tue, 20 Oct 2026 09:01:47 +0200 Subject: [PATCH 13/13] fix(tracing): preserve raw caller context values Signed-off-by: rtemperini <29879569+rtemperini@users.noreply.github.com> --- go/adk/pkg/telemetry/context_attributes.go | 4 +- .../pkg/telemetry/context_attributes_test.go | 70 ++++++++++++++++++- .../tests/test_caller_context_attributes.py | 42 ++++++++++- 3 files changed, 112 insertions(+), 4 deletions(-) diff --git a/go/adk/pkg/telemetry/context_attributes.go b/go/adk/pkg/telemetry/context_attributes.go index 6bd8923296..72e09c9686 100644 --- a/go/adk/pkg/telemetry/context_attributes.go +++ b/go/adk/pkg/telemetry/context_attributes.go @@ -72,7 +72,9 @@ func ContextWithPromotedMetadata(ctx context.Context, metadata map[string]any) c if value == "" { continue } - member, err := baggage.NewMember(mapping.attribute, value) + // Metadata and decoded baggage already hold raw values. The + // propagator handles header encoding when crossing a process boundary. + member, err := baggage.NewMemberRaw(mapping.attribute, value) if err != nil { continue } diff --git a/go/adk/pkg/telemetry/context_attributes_test.go b/go/adk/pkg/telemetry/context_attributes_test.go index fd687466a9..1392da6683 100644 --- a/go/adk/pkg/telemetry/context_attributes_test.go +++ b/go/adk/pkg/telemetry/context_attributes_test.go @@ -9,6 +9,7 @@ import ( "go.opentelemetry.io/contrib/processors/baggagecopy" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/baggage" + "go.opentelemetry.io/otel/propagation" sdktrace "go.opentelemetry.io/otel/sdk/trace" "go.opentelemetry.io/otel/sdk/trace/tracetest" ) @@ -25,9 +26,9 @@ func baggageContext(t *testing.T, members map[string]string) context.Context { built := make([]baggage.Member, 0, len(members)) for key, value := range members { - member, err := baggage.NewMember(key, value) + member, err := baggage.NewMemberRaw(key, value) if err != nil { - t.Fatalf("baggage.NewMember(%q, %q): %v", key, value, err) + t.Fatalf("baggage.NewMemberRaw(%q, %q): %v", key, value, err) } built = append(built, member) } @@ -298,3 +299,68 @@ func mustMember(t *testing.T, key, value string) baggage.Member { } return member } + +func TestPromotedContextPreservesRawValues(t *testing.T) { + values := []struct{ name, value string }{ + {"spaces", "case with spaces"}, + {"unicode", "café"}, + {"percent escape", "literal%2Fid"}, + {"double percent escape", "literal%252Fid"}, + {"literal percent", "100%"}, + {"delimiters", "a,b;c=d"}, + } + for _, source := range []string{"metadata", "mapped metadata", "baggage"} { + for _, value := range values { + t.Run(source+"/"+value.name, func(t *testing.T) { + policy := "thread_id" + key := "thread_id" + if source != "metadata" { + policy = `[{"from":"source","to":"thread_id"}]` + key = "source" + } + setAllowlist(t, policy) + ctx := t.Context() + metadata := map[string]any{key: value.value, "private_token": "must-not-be-exported"} + if source == "baggage" { + ctx = baggageContext(t, map[string]string{key: value.value}) + metadata = nil + } + ctx = ContextWithPromotedMetadata(ctx, metadata) + if got := baggage.FromContext(ctx).Member("thread_id").Value(); got != value.value { + t.Fatalf("promoted value = %q, want %q", got, value.value) + } + + exporter := tracetest.NewInMemoryExporter() + provider := sdktrace.NewTracerProvider( + sdktrace.WithSyncer(exporter), + sdktrace.WithSpanProcessor(baggagecopy.NewSpanProcessor(AllowedBaggageCopyFilter())), + ) + t.Cleanup(func() { _ = provider.Shutdown(context.Background()) }) + tracer := provider.Tracer("test") + ctx, root := tracer.Start(ctx, "invocation") + _, model := tracer.Start(ctx, "model") + model.End() + root.End() + for _, name := range []string{"invocation", "model"} { + attrs := spanAttributesByName(t, exporter.GetSpans(), name) + if got := attrs["thread_id"].AsString(); got != value.value { + t.Errorf("%s thread_id = %q, want %q", name, got, value.value) + } + if _, exists := attrs["private_token"]; exists { + t.Errorf("%s exported unrelated metadata", name) + } + } + + // Encoding belongs to the propagator. A downstream hop must + // receive the original value, not a decoded percent sequence. + carrier := propagation.MapCarrier{} + propagator := propagation.Baggage{} + propagator.Inject(ctx, carrier) + downstream := propagator.Extract(t.Context(), carrier) + if got := baggage.FromContext(downstream).Member("thread_id").Value(); got != value.value { + t.Errorf("propagated value = %q, want %q", got, value.value) + } + }) + } + } +} diff --git a/python/packages/kagent-core/tests/test_caller_context_attributes.py b/python/packages/kagent-core/tests/test_caller_context_attributes.py index 6c75c0c7e3..5c57531d6d 100644 --- a/python/packages/kagent-core/tests/test_caller_context_attributes.py +++ b/python/packages/kagent-core/tests/test_caller_context_attributes.py @@ -3,8 +3,9 @@ import pytest from a2a.types import Message, Role from google.protobuf.json_format import ParseDict -from opentelemetry import baggage +from opentelemetry import baggage, trace from opentelemetry import context as otel_context +from opentelemetry.baggage.propagation import W3CBaggagePropagator from opentelemetry.processor.baggage import BaggageSpanProcessor from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor @@ -262,3 +263,42 @@ def test_promotes_real_protobuf_message_metadata(monkeypatch): "user.id": "opaque-subject", "attempt": "3", } + + +@pytest.mark.parametrize("value", ["case with spaces", "café", "literal%2Fid", "literal%252Fid", "100%", "a,b;c=d"]) +@pytest.mark.parametrize("source", ["metadata", "mapped metadata", "baggage"]) +def test_promoted_context_preserves_raw_values(monkeypatch, value, source): + policy = "thread_id" if source == "metadata" else '[{"from":"source","to":"thread_id"}]' + key = "thread_id" if source == "metadata" else "source" + set_allowlist(monkeypatch, policy) + ctx = otel_context.Context() + metadata = {key: value, "private_token": "must-not-be-exported"} + if source == "baggage": + ctx = baggage_context({key: value}) + metadata = None + ctx = context_with_promoted_metadata(metadata, ctx) + assert baggage.get_baggage("thread_id", ctx) == value + + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(BaggageSpanProcessor(allowed_baggage_key_predicate())) + provider.add_span_processor(SimpleSpanProcessor(exporter)) + try: + tracer = provider.get_tracer("test") + root = tracer.start_span("invocation", context=ctx) + child_ctx = trace.set_span_in_context(root, ctx) + model = tracer.start_span("model", context=child_ctx) + model.end() + root.end() + for name in ("invocation", "model"): + attrs = _span_attrs(exporter.get_finished_spans(), name) + assert attrs["thread_id"] == value + assert "private_token" not in attrs + + carrier = {} + propagator = W3CBaggagePropagator() + propagator.inject(carrier, context=ctx) + downstream = propagator.extract(carrier, context=otel_context.Context()) + assert baggage.get_baggage("thread_id", downstream) == value + finally: + provider.shutdown()