diff --git a/docs/architecture/README.md b/docs/architecture/README.md index a0de4fd65..89a6376c0 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 000000000..aed95a69a --- /dev/null +++ b/docs/architecture/trace-context.md @@ -0,0 +1,175 @@ +# 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 — +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. + +--- + +## Configuration + +| Setting | Default | Description | +|---|---|---| +| 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: + - {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. + +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 + +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" +``` + +| 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 +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. + +--- + +## 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 +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 | 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 +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 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["session.id"], span.attributes["kagent.thread_id"]) + where span.attributes["kagent.thread_id"] != 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. diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index 57d312007..220180648 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 000000000..d0774d144 --- /dev/null +++ b/go/adk/pkg/telemetry/context_attributes.go @@ -0,0 +1,272 @@ +package telemetry + +import ( + "context" + "crypto/hmac" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "os" + "strconv" + "strings" + "unicode" + + "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. + 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 +// 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 +// 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 { + mappings := allowedContextMappings() + if len(mappings) == 0 { + return nil + } + + bag := baggage.FromContext(ctx) + 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 + } + 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 + } + 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 { + raw := strings.TrimSpace(os.Getenv(traceContextKeysEnvVar)) + if raw == "" { + return nil + } + if strings.HasPrefix(raw, "[") { + return capMappings(parseJSONAllowlist(raw)) + } + return capMappings(parseCommaAllowlist(raw)) +} + +func parseCommaAllowlist(raw string) []contextMapping { + mappings := make([]contextMapping, 0, maxContextKeys) + for key := range strings.SplitSeq(raw, ",") { + 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 + } + 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[id] = struct{}{} + out = append(out, mapping) + if len(out) == maxContextKeys { + break + } + } + 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. +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 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) { + return -1 + } + return r + }, 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 new file mode 100644 index 000000000..89d996218 --- /dev/null +++ b/go/adk/pkg/telemetry/context_attributes_test.go @@ -0,0 +1,358 @@ +package telemetry + +import ( + "context" + "crypto/hmac" + "crypto/sha256" + "encoding/hex" + "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 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: "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: "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: "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", "extra": "nope"}, + 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"}, + }, + { + // 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) + + 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_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") + + 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 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(allowedContextMappings()); got != maxContextKeys { + t.Errorf("allowlist length = %d, want %d", got, maxContextKeys) + } +} + +func TestAllowedContextMappings_DropsOverLongAndDuplicateKeys(t *testing.T) { + t.Setenv(traceContextKeysEnvVar, "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) + } + for i, key := range want { + if got[i].source != key { + t.Errorf("key %d = %q, want %q", i, got[i].source, 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, `[{"from":"sub","to":"user.id"}]`) + + 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{"sub": "opaque-subject"}) + 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["user.id"].AsString(); got != "opaque-subject" { + t.Errorf("span %q: user.id = %q, want %q", name, got, "opaque-subject") + } + } +} + +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{"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 { + key := string(attr.Key) + if strings.HasPrefix(key, contextAttributePrefix) || key == "user.id" { + t.Errorf("unexpected promoted attribute %q", attr.Key) + } + } +} + +// 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") + + 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) + } +} diff --git a/go/core/pkg/env/otel.go b/go/core/pkg/env/otel.go index cf3dc1b69..2b8775b65 100644 --- a/go/core/pkg/env/otel.go +++ b/go/core/pkg/env/otel.go @@ -37,4 +37,23 @@ var ( "OTLP exporter endpoint for logs. Takes precedence over OTEL_EXPORTER_OTLP_ENDPOINT for logs.", ComponentController, ) + + 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.*, 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 771bff0ca..e23edcbc4 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,94 @@ 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 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"}}, + } + harnessSupplied := "tenant.supplied" + + 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: "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", + keys: "sub", + harnessEnv: []v1alpha3.HarnessEnvVar{{Name: env.KagentTraceContextKeys.Name(), Value: &harnessSupplied}}, + 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.keys) + t.Setenv(env.KagentTraceContextHashKey.Name(), tt.hashKey) + 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) + + gotKeys, seenKeys := "", 0 + gotHash, seenHash := "", 0 + for _, variable := range revision.Environment { + 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, 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) + }) + } +} + 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 5be6d5890..7373b7cae 100644 --- a/go/core/v2/translator/kagent/compiler.go +++ b/go/core/v2/translator/kagent/compiler.go @@ -91,6 +91,13 @@ 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. 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 @@ -432,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 044859777..02c0468e2 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 d31b3e845..c4a93c89d 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: {{ include "kagent.traceContextKeys" . | 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/templates/controller-deployment.yaml b/helm/kagent/templates/controller-deployment.yaml index 7e74c9437..d111719dd 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 a30e1619c..b2f3ad4b5 100644 --- a/helm/kagent/tests/controller-deployment_test.yaml +++ b/helm/kagent/tests/controller-deployment_test.yaml @@ -831,3 +831,59 @@ 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: + - sub + - thread_id + - channel + asserts: + - equal: + path: data.KAGENT_TRACE_CONTEXT_KEYS + 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 daa3fe9f8..31942c130 100644 --- a/helm/kagent/values.yaml +++ b/helm/kagent/values.yaml @@ -848,6 +848,31 @@ 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. + # + # 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-adk/src/kagent/adk/_agent_executor.py b/python/packages/kagent-adk/src/kagent/adk/_agent_executor.py index 18107e1ca..fb8cf98c5 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 73cc2a650..4fafa5409 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 74608dbd8..568d71d19 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 826775371..87307f1e7 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 000000000..67853b394 --- /dev/null +++ b/python/packages/kagent-core/src/kagent/core/tracing/_context_attributes.py @@ -0,0 +1,251 @@ +"""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 + +# 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" + +# 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, +) -> 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 + 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: + Attribute name to sanitised value. Empty when the allowlist is empty, + which is the default. + """ + mappings = _allowed_context_mappings() + if not mappings: + return {} + + bag = baggage.get_all(context) + attributes: dict[str, str] = {} + 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: + value = _sanitize_context_value(scalar) + if not value: + 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: + continue + attributes[name] = value + return attributes + + +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 + 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 [] + if raw.startswith("["): + return _cap_mappings(_parse_json_allowlist(raw)) + return _cap_mappings(_parse_comma_allowlist(raw)) + + +def _parse_comma_allowlist(raw: str) -> list[_ContextMapping]: + mappings: list[_ContextMapping] = [] + for candidate in raw.split(","): + 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 + seen.add(identity) + out.append(mapping) + if len(out) == MAX_CONTEXT_KEYS: + break + 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: + """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: + """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() + + +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/tests/test_caller_context_attributes.py b/python/packages/kagent-core/tests/test_caller_context_attributes.py new file mode 100644 index 000000000..2584f84cc --- /dev/null +++ b/python/packages/kagent-core/tests/test_caller_context_attributes.py @@ -0,0 +1,249 @@ +import hashlib +import hmac + +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_HASH_KEY_ENV_VAR, + TRACE_CONTEXT_KEYS_ENV_VAR, + _allowed_context_mappings, +) +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 + + +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_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", + } + + 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, "sub") + assert caller_context_attributes( + {"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", "extra": "nope"}, + 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): + """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 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_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_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() == [] + + +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, '[{"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)] 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 4abd6cbfa..35112a4a0 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 3ad0b3e3f..b1415c66d 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": "sub=opaque-subject"}) + + assert get_baggage("sub", ctx) == "opaque-subject" + + @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 001024772..dc5a80b10 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 f3ff8a6f8..46c76ddc3 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