diff --git a/docs/architecture/README.md b/docs/architecture/README.md index 4ea4f4c487..e7815806fb 100644 --- a/docs/architecture/README.md +++ b/docs/architecture/README.md @@ -113,6 +113,7 @@ cleanup. Session lifecycle workers independently pause/suspend idle Actors. - [Token usage](usage-extension.md) - [Prompt resolution](prompt-templates.md) - [Telemetry](telemetry.md) +- [Caller context in traces](trace-context.md) - [Structured output](structured-output.md) The documents describe implemented behavior. Full cross-Session delegation diff --git a/docs/architecture/trace-context.md b/docs/architecture/trace-context.md new file mode 100644 index 0000000000..b579e884aa --- /dev/null +++ b/docs/architecture/trace-context.md @@ -0,0 +1,154 @@ +# Caller Context in Traces + +Agent spans describe *what the agent did*, but they say nothing about *who asked +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. + +The feature is **off by default**. It turns on when an operator sets an +allowlist. + +--- + +## Configuration + +| Setting | Default | Description | +|---|---|---| +| Helm `controller.traceContextKeys` | `[]` | 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 +controller: + traceContextKeys: + - {from: sub, to: user.id} + - thread_id + - gen_ai.conversation.id +``` + +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 + +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. + +A2A `message.metadata` is promoted *into baggage* at the request 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 (`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 request boundary 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. Metadata is promoted +before the invocation span starts, so the allowlist also covers that span. + +Baggage propagation across processes is separately opt-in on current kagent. +Set `OTEL_PROPAGATORS=tracecontext,baggage` on the controller and each runtime +that should exchange baggage. Use `controller.env` for the controller and the +Harness workload environment for runtimes. Controller `OTEL_*` overrides are +process-local; they are not forwarded to managed runtimes. Metadata promotion +and local span attributes do not require outbound baggage propagation. + +--- + +## Where values land + +``` +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" +``` + +The request-scoped kagent attribute bag still carries runtime values +(`enduser.id`, `gen_ai.conversation.id`, `a2a.task.id`, …). That processor runs after the +baggage processor, so caller baggage cannot override attributes the runtime already +stamped. + +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. + +--- + +## 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. + +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. + +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: + +| Risk | Control | +|---|---| +| 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 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. + +--- + +## Renaming attributes for a 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["thread_id"]) + where span.attributes["thread_id"] != nil +``` + +--- + +## Implementation + +| Component | Path | +|---|---| +| Go ADK | `go/adk/pkg/telemetry/context_attributes.go` | +| Python core | `python/packages/kagent-core/src/kagent/core/telemetry/_context_attributes.py` | +| Python A2A boundary (all runtimes) | `python/packages/kagent-core/src/kagent/core/a2a/_telemetry.py` | +| Controller forwarding | `go/core/internal/translator/kagent/compiler.go` | +| Helm | `helm/kagent/templates/controller-deployment.yaml` | diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index dd6469b7cf..d24083e9dc 100644 --- a/go/adk/pkg/a2a/executor.go +++ b/go/adk/pkg/a2a/executor.go @@ -275,10 +275,12 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.ExecutorCon defer e.runningUsage.track(reqCtx.TaskID, usage)() // The invocation span started before this executor ran, so the request // identity has to be recorded on it directly. ADK's own spans get it - // through the request attribute span processor. + // through the request attribute span processor. Promote here as well + // for invocations that bypass the transport interceptor. resumed := reqCtx.StoredTask != nil && (reqCtx.StoredTask.Status.State == a2atype.TaskStateInputRequired || reqCtx.StoredTask.Status.State == a2atype.TaskStateAuthRequired) 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)...) // The identity replaces any caller members of the same name, so a caller // cannot impersonate another conversation downstream. diff --git a/go/adk/pkg/a2a/executor_test.go b/go/adk/pkg/a2a/executor_test.go index 3039db4909..46eb009cca 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/a2a/server/tracing.go b/go/adk/pkg/a2a/server/tracing.go index a04e53f0ab..654400832d 100644 --- a/go/adk/pkg/a2a/server/tracing.go +++ b/go/adk/pkg/a2a/server/tracing.go @@ -11,6 +11,7 @@ import ( "go.opentelemetry.io/otel/metric" "github.com/kagent-dev/kagent/go/adk/pkg/auth" + "github.com/kagent-dev/kagent/go/adk/pkg/telemetry" "github.com/kagent-dev/kagent/go/pkg/tracing" ) @@ -71,11 +72,20 @@ func newInvocationInterceptor(logger *slog.Logger, telemetry tracing.RuntimeTele return interceptor } -func (i *invocationInterceptor) Before(ctx context.Context, callCtx *a2asrv.CallContext, _ *a2asrv.Request) (context.Context, any, error) { +func (i *invocationInterceptor) Before(ctx context.Context, callCtx *a2asrv.CallContext, request *a2asrv.Request) (context.Context, any, error) { method := callCtx.Method() if method != "SendMessage" && method != "SendStreamingMessage" { return ctx, nil, nil } + var metadata map[string]any + if request != nil { + if params, ok := request.Payload.(*a2atype.SendMessageRequest); ok && params != nil && params.Message != nil { + metadata = params.Message.Metadata + } + } + // Promote before the invocation starts, so the same baggage processor + // covers the invocation and every descendant span. + ctx = telemetry.ContextWithPromotedMetadata(ctx, metadata) attributes := make([]attribute.KeyValue, 0, len(i.static)+2) attributes = append(attributes, attribute.String(tracing.AttributeMethod, method)) attributes = append(attributes, i.static...) @@ -85,6 +95,8 @@ func (i *invocationInterceptor) Before(ctx context.Context, callCtx *a2asrv.Call attributes = append(attributes, attribute.String(tracing.AttributeUserID, userID)) } ctx, invocation := tracing.StartInvocation(ctx, tracing.Tracer(invocationScope), i.name, i.flush, attributes...) + // Runtime identity remains authoritative after baggagecopy runs. + invocation.SetAttributes(attributes...) invocation.MeasureDuration(i.duration, i.durationAttributes...) // a2a-go runs no final After callback when a streaming consumer stops // reading, so an invocation the transport still owns would otherwise never diff --git a/go/adk/pkg/a2a/server/tracing_test.go b/go/adk/pkg/a2a/server/tracing_test.go index 20b69d1d60..8ee86461e4 100644 --- a/go/adk/pkg/a2a/server/tracing_test.go +++ b/go/adk/pkg/a2a/server/tracing_test.go @@ -11,6 +11,9 @@ import ( "log/slog" "net/http" "net/http/httptest" + "os" + "os/exec" + "strings" "testing" "time" @@ -29,7 +32,9 @@ import ( "google.golang.org/grpc/credentials/insecure" "github.com/kagent-dev/kagent/go/adk/pkg/auth" + "github.com/kagent-dev/kagent/go/adk/pkg/telemetry" "github.com/kagent-dev/kagent/go/pkg/tracing" + "go.opentelemetry.io/contrib/processors/baggagecopy" ) // staticUserInterceptor stands in for UserIDCallInterceptor, which the app @@ -501,3 +506,57 @@ func TestRequestSpanCompletesOnAStreamingMessageResult(t *testing.T) { requestSpan(t, exporter) } + +func TestInvocationPromotesOnlyAllowlistedMetadata(t *testing.T) { + mode := os.Getenv("KAGENT_CONTEXT_POLICY_TEST") + if mode == "" { + // Policy is cached for process lifetime, so each mode needs a fresh process. + for _, mode := range []string{"disabled", "enabled"} { + t.Run(mode, func(t *testing.T) { + policy := "" + if mode == "enabled" { + policy = `[{"from":"sub","to":"user.id"},"gen_ai.agent.name","enduser.id"]` + } + command := exec.CommandContext(t.Context(), os.Args[0], "-test.run=^TestInvocationPromotesOnlyAllowlistedMetadata$") + command.Env = append(os.Environ(), "KAGENT_CONTEXT_POLICY_TEST="+mode, "KAGENT_TRACE_CONTEXT_KEYS="+policy) + if output, err := command.CombinedOutput(); err != nil { + t.Fatalf("%s policy: %v\n%s", mode, err, output) + } + }) + } + return + } + exporter := syncExporter(t) + if filter := telemetry.AllowedBaggageCopyFilter(); filter != nil { + otel.GetTracerProvider().(*sdktrace.TracerProvider).RegisterSpanProcessor(baggagecopy.NewSpanProcessor(filter)) + } + server, err := NewA2AServer(a2atype.AgentCard{}, failingExecutor{}, slog.New(slog.DiscardHandler), + ServerConfig{Port: "0", Telemetry: tracing.RuntimeTelemetry{ + Runtime: tracing.RuntimeADKGo, AgentName: "researcher", AgentNamespace: "team-a", + }}, a2asrv.WithCallInterceptors(staticUserInterceptor{userID: "trusted-user"})) + if err != nil { + t.Fatal(err) + } + message := a2atype.NewMessage(a2atype.MessageRoleUser, a2atype.NewTextPart("hello")) + message.Metadata = map[string]any{ + "sub": "opaque-subject", "private_token": "must-not-be-exported", + "gen_ai.agent.name": "forged-agent", "enduser.id": "forged-user", + } + sendMessage(t, server, &a2atype.SendMessageRequest{Message: message}) + span := requestSpan(t, exporter) + wantUser := "" + if mode == "enabled" { + wantUser = "opaque-subject" + } + if got := spanAttribute(span, "user.id"); got != wantUser { + t.Fatalf("user.id = %q, want %q", got, wantUser) + } + if spanAttribute(span, tracing.AttributeAgentName) != "researcher" || spanAttribute(span, tracing.AttributeUserID) != "trusted-user" { + t.Fatal("caller metadata replaced runtime identity") + } + for _, attr := range span.Attributes { + if strings.HasPrefix(string(attr.Key), "a2a.message.metadata.") || string(attr.Key) == "private_token" || attr.Value.AsString() == "must-not-be-exported" { + t.Errorf("unallowlisted metadata exported on invocation: %s", attr.Key) + } + } +} diff --git a/go/adk/pkg/telemetry/attributes_test.go b/go/adk/pkg/telemetry/attributes_test.go index def6219815..0ec31e6f1a 100644 --- a/go/adk/pkg/telemetry/attributes_test.go +++ b/go/adk/pkg/telemetry/attributes_test.go @@ -48,3 +48,20 @@ func TestWithRequestAttributesKeepsContextWhenEmpty(t *testing.T) { t.Fatal("no attributes should return the same context") } } +func spanAttributesByName(t *testing.T, spans tracetest.SpanStubs, name string) map[string]attribute.Value { + t.Helper() + + for _, span := range spans { + if span.Name != name { + continue + } + attrs := make(map[string]attribute.Value, len(span.Attributes)) + for _, attr := range span.Attributes { + attrs[string(attr.Key)] = attr.Value + } + return attrs + } + + t.Fatalf("span %q not found", name) + return nil +} diff --git a/go/adk/pkg/telemetry/context_attributes.go b/go/adk/pkg/telemetry/context_attributes.go new file mode 100644 index 0000000000..72e09c9686 --- /dev/null +++ b/go/adk/pkg/telemetry/context_attributes.go @@ -0,0 +1,232 @@ +package telemetry + +import ( + "context" + "encoding/json" + "os" + "strconv" + "strings" + "sync" + "unicode" + + "github.com/kagent-dev/kagent/go/core/pkg/env" + "go.opentelemetry.io/contrib/processors/baggagecopy" + "go.opentelemetry.io/otel/baggage" +) + +const ( + maxContextKeys = 32 + maxContextKeyLength = 64 +) + +// 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 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 +} + +// loadContextMappings parses the allowlist once per process. Tests reset it +// after changing the environment. +var loadContextMappings = sync.OnceValue(parseAllowedContextMappings) + +func resetAllowedContextMappings() { + loadContextMappings = sync.OnceValue(parseAllowedContextMappings) +} + +func allowedContextMappings() []contextMapping { + return loadContextMappings() +} + +// 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. +// +// 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 ctx + } + + bag := baggage.FromContext(ctx) + changed := false + for _, mapping := range mappings { + if bag.Member(mapping.attribute).Value() != "" { + continue + } + value := "" + if scalar, ok := scalarString(metadata[mapping.source]); ok { + value = sanitizeContextValue(scalar) + } + if value == "" && mapping.source != mapping.attribute { + value = sanitizeContextValue(bag.Member(mapping.source).Value()) + } + if value == "" { + continue + } + // 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 + } + next, err := bag.SetMember(member) + if err != nil { + continue + } + bag = next + changed = true + } + if !changed { + return ctx + } + return baggage.ContextWithBaggage(ctx, bag) +} + +// 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 + } +} + +func parseAllowedContextMappings() []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 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) + continue + } + } + return mappings +} + +func newContextMapping(from, to string) (contextMapping, bool) { + from = strings.TrimSpace(from) + to = strings.TrimSpace(to) + if from == "" || len([]rune(from)) > maxContextKeyLength || !isAttributeKey(from) { + return contextMapping{}, false + } + if to == "" { + to = from + } + if len([]rune(to)) > maxContextKeyLength || !isAttributeKey(to) { + return contextMapping{}, false + } + 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 + if _, duplicate := seen[id]; duplicate { + continue + } + seen[id] = struct{}{} + out = append(out, mapping) + if len(out) == maxContextKeys { + break + } + } + return out +} + +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 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: + return v, true + case bool: + return strconv.FormatBool(v), true + case float64: + return strconv.FormatFloat(v, 'f', -1, 64), 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) +} 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..1392da6683 --- /dev/null +++ b/go/adk/pkg/telemetry/context_attributes_test.go @@ -0,0 +1,366 @@ +package telemetry + +import ( + "context" + "strconv" + "strings" + "testing" + + "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" +) + +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() + + built := make([]baggage.Member, 0, len(members)) + for key, value := range members { + member, err := baggage.NewMemberRaw(key, value) + if err != nil { + t.Fatalf("baggage.NewMemberRaw(%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 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 TestContextWithPromotedMetadata(t *testing.T) { + tests := []struct { + name string + allowlist string + baggageVals map[string]string + metadata map[string]any + want map[string]string + }{ + { + 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: "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: "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 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: "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: "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: "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: "skips non-scalars", + allowlist: "value", + metadata: map[string]any{"value": map[string]any{"nested": true}}, + want: map[string]string{}, + }, + { + 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{"note": "linebreakandnul"}, + }, + { + 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"}, + }, + { + 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 { + t.Run(tt.name, func(t *testing.T) { + setAllowlist(t, tt.allowlist) + got := baggageValues(ContextWithPromotedMetadata(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: got %q, want %q", key, got[key], 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" + } + 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 TestAllowedBaggageCopyFilter(t *testing.T) { + if AllowedBaggageCopyFilter() != nil { + t.Fatal("empty allowlist must not install a baggagecopy filter") + } + + setAllowlist(t, `[{"from":"sub","to":"user.id"},"thread_id"]`) + filter := AllowedBaggageCopyFilter() + if filter == nil { + t.Fatal("expected a filter") + } + if !filter(mustMember(t, "user.id", "x")) { + t.Error("user.id should be copied") + } + if !filter(mustMember(t, "thread_id", "x")) { + t.Error("thread_id should be copied") + } + if filter(mustMember(t, "sub", "x")) { + t.Error("source-only key sub should not be copied") + } + if filter(mustMember(t, "secret", "x")) { + t.Error("non-allowlisted key should not be copied") + } +} + +func TestCallerContextLandsOnEverySpan(t *testing.T) { + setAllowlist(t, "user.id") + + exporter := tracetest.NewInMemoryExporter() + tp := sdktrace.NewTracerProvider( + sdktrace.WithSyncer(exporter), + sdktrace.WithSpanProcessor(baggagecopy.NewSpanProcessor(AllowedBaggageCopyFilter())), + sdktrace.WithSpanProcessor(requestAttributesSpanProcessor{}), + ) + t.Cleanup(func() { _ = tp.Shutdown(context.Background()) }) + + ctx := ContextWithPromotedMetadata(baggageContext(t, map[string]string{"user.id": "opaque-subject"}), nil) + ctx = WithRequestAttributes(ctx, attribute.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() + + 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) + } + } +} + +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(requestAttributesSpanProcessor{}), + ) + t.Cleanup(func() { _ = tp.Shutdown(context.Background()) }) + + ctx := ContextWithPromotedMetadata(context.Background(), map[string]any{"sub": "opaque-subject"}) + ctx = WithRequestAttributes(ctx, attribute.String("gen_ai.conversation.id", "sess-runtime")) + tracer := tp.Tracer("test") + ctx, root := tracer.Start(ctx, "invocation") + _, model := tracer.Start(ctx, "generate_content") + model.End() + root.End() + + for _, name := range []string{"invocation", "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["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 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(requestAttributesSpanProcessor{}), + ) + t.Cleanup(func() { _ = tp.Shutdown(context.Background()) }) + + ctx := baggageContext(t, map[string]string{"gen_ai.conversation.id": "from-caller"}) + ctx = WithRequestAttributes(ctx, attribute.String("gen_ai.conversation.id", "from-runtime")) + tracer := tp.Tracer("test") + _, span := tracer.Start(ctx, "generate_content") + span.End() + + 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) + } +} + +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 +} + +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/go/adk/pkg/telemetry/tracing.go b/go/adk/pkg/telemetry/tracing.go index 88c24f0bd3..08b4fcd959 100644 --- a/go/adk/pkg/telemetry/tracing.go +++ b/go/adk/pkg/telemetry/tracing.go @@ -3,18 +3,24 @@ package telemetry import ( "context" + "go.opentelemetry.io/contrib/processors/baggagecopy" sdktrace "go.opentelemetry.io/otel/sdk/trace" kagenttelemetry "github.com/kagent-dev/kagent/go/pkg/telemetry" "github.com/kagent-dev/kagent/go/pkg/tracing" ) -// Init boots the process telemetry with the request attribute span processor. +// Init installs caller baggage promotion before the authoritative request attributes. func Init(ctx context.Context, telemetry tracing.RuntimeTelemetry) (*kagenttelemetry.Providers, error) { + processors := make([]sdktrace.SpanProcessor, 0, 2) + if filter := AllowedBaggageCopyFilter(); filter != nil { + processors = append(processors, baggagecopy.NewSpanProcessor(filter)) + } + processors = append(processors, requestAttributesSpanProcessor{}) return kagenttelemetry.Init(ctx, kagenttelemetry.Options{ Identity: telemetry.ResourceIdentity(), Config: telemetry.ExporterConfig(), Defaults: telemetry.ResourceDefaults(""), - SpanProcessors: []sdktrace.SpanProcessor{requestAttributesSpanProcessor{}}, + SpanProcessors: processors, }) } diff --git a/go/core/internal/translator/kagent/compiler.go b/go/core/internal/translator/kagent/compiler.go index 3593770b7e..5f16270b53 100644 --- a/go/core/internal/translator/kagent/compiler.go +++ b/go/core/internal/translator/kagent/compiler.go @@ -10,12 +10,14 @@ import ( a2atype "github.com/a2aproject/a2a-go/v2/a2a" "github.com/a2aproject/a2a-go/v2/a2apb/v1/pbconv" "istio.io/istio/pkg/kube/krt" + corev1 "k8s.io/api/core/v1" runtimeconfig "github.com/kagent-dev/kagent/go/adk/pkg/config" apia2a "github.com/kagent-dev/kagent/go/api/a2a" configv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/config/v1" "github.com/kagent-dev/kagent/go/core/internal/translator" "github.com/kagent-dev/kagent/go/core/internal/translator/adkconfig" + "github.com/kagent-dev/kagent/go/core/pkg/env" ) // Compiler translates resolved inputs into a kagent runtime revision. @@ -58,6 +60,8 @@ func (c *Compiler) Compile(ctx context.Context, input *translator.HarnessInput) harnessEnvironment := adkconfig.HarnessEnvironment(harness) environment := append(compiled.Environment, harnessEnvironment...) + // Caller context policy is operator-owned; a Harness cannot enable or widen it. + environment = applyOperatorOnlyEnv(environment, env.KagentTraceContextKeys) environment = adkconfig.DedupeEnv(environment) provenance, err := c.config.BuildProvenance(ctx, harness, compiled.Templates, compiled.Models, environment) if err != nil { @@ -88,3 +92,12 @@ func agentCard(agentName string, template *translator.TemplateConfiguration) *a2 card.Capabilities.Extensions = append(card.Capabilities.Extensions, apia2a.UsageExtension()) return card } + +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 +} 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/go/core/pkg/env/otel.go b/go/core/pkg/env/otel.go index 0a228f0798..60938eec19 100644 --- a/go/core/pkg/env/otel.go +++ b/go/core/pkg/env/otel.go @@ -12,6 +12,8 @@ type OTELSignalVars struct { } var ( + KagentTraceContextKeys = RegisterStringVar("KAGENT_TRACE_CONTEXT_KEYS", "", + "Allowlist of caller context keys copied from W3C baggage onto agent spans. Accepts comma-separated keys or JSON strings and {from, to} mappings. Empty disables promotion.", ComponentController, ComponentAgentRuntime) OTELTraces = registerOTELSignal("TRACES") OTELMetrics = registerOTELSignal("METRICS") OTELLogs = registerOTELSignal("LOGS") diff --git a/go/go.mod b/go/go.mod index e6d9e82690..29e065341c 100644 --- a/go/go.mod +++ b/go/go.mod @@ -54,6 +54,7 @@ require ( go.opentelemetry.io/contrib/exporters/autoexport v0.72.0 go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.72.0 go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.72.0 + go.opentelemetry.io/contrib/processors/baggagecopy v0.16.3 go.opentelemetry.io/contrib/propagators/autoprop v0.72.0 go.opentelemetry.io/otel v1.47.0 go.opentelemetry.io/otel/exporters/prometheus v0.69.0 diff --git a/go/go.sum b/go/go.sum index 20c4cb4a16..f25125c526 100644 --- a/go/go.sum +++ b/go/go.sum @@ -1231,6 +1231,8 @@ go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.7 go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.72.0/go.mod h1:A+EVhDakAj1waJBPhmTszZaSr08ww8MRRKy1nB1rqck= go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.72.0 h1:LxwW/9ctSCv+QkE/cLR7M91ZIkXNMqJtEMi1vCw9U8s= go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.72.0/go.mod h1:tOsftB4SslBwwErVEPaenU2RpThXWPIU8DoJHEC4dyw= +go.opentelemetry.io/contrib/processors/baggagecopy v0.16.3 h1:rTjkMS+1AKIHd51/nmZWhvk/fq2WQ+POpA+TRt1+q7k= +go.opentelemetry.io/contrib/processors/baggagecopy v0.16.3/go.mod h1:dFS4IiUCeRVkpd83JqPK8X06gUJXWto/APtRE/mC4ew= go.opentelemetry.io/contrib/propagators/autoprop v0.72.0 h1:/Kv8DucXan87cGRpa3NeZtlk26kFahF2YrL6Sz3GFsI= go.opentelemetry.io/contrib/propagators/autoprop v0.72.0/go.mod h1:Qttdialu54uM051nRjcHgqPshuTi2YGZb/rruiNJ8zo= go.opentelemetry.io/contrib/propagators/aws v1.47.0 h1:n/2x/vQki1ojD7R2jYHdnNdI74Yx/OFdY4YXFFesm8M= diff --git a/helm/kagent/templates/_helpers.tpl b/helm/kagent/templates/_helpers.tpl index ca3d36d90b..1b898ddabd 100644 --- a/helm/kagent/templates/_helpers.tpl +++ b/helm/kagent/templates/_helpers.tpl @@ -438,3 +438,20 @@ tag wins over the component tag. {{- define "kagent.database.skipMigrations" -}} {{- if dig "database" "skipMigrations" false (.Values.config | default dict) -}}true{{- end -}} {{- end -}} + +{{/* +Serialize controller.traceContextKeys for KAGENT_TRACE_CONTEXT_KEYS. +A list of strings is joined with commas; any mapping entry is emitted as JSON +so {from, to} 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-deployment.yaml b/helm/kagent/templates/controller-deployment.yaml index 824f44efcd..4bc5d4942b 100644 --- a/helm/kagent/templates/controller-deployment.yaml +++ b/helm/kagent/templates/controller-deployment.yaml @@ -150,6 +150,10 @@ spec: valueFrom: fieldRef: fieldPath: spec.nodeName + {{- with .Values.controller.traceContextKeys }} + - name: KAGENT_TRACE_CONTEXT_KEYS + value: {{ include "kagent.traceContextKeys" . | quote }} + {{- end }} {{- with .Values.controller.env }} {{- toYaml . | nindent 12 }} {{- end }} diff --git a/helm/kagent/tests/trace-context_test.yaml b/helm/kagent/tests/trace-context_test.yaml new file mode 100644 index 0000000000..b41cb79941 --- /dev/null +++ b/helm/kagent/tests/trace-context_test.yaml @@ -0,0 +1,36 @@ +suite: caller trace context policy +templates: + - controller-deployment.yaml + - controller-configmap.yaml +tests: + - it: should leave caller context disabled by default + template: controller-deployment.yaml + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: KAGENT_TRACE_CONTEXT_KEYS + any: true + - it: should join string keys into the controller policy + set: + controller.traceContextKeys: [sub, thread_id, channel] + template: controller-deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: KAGENT_TRACE_CONTEXT_KEYS + value: sub,thread_id,channel + - it: should preserve mappings as JSON + set: + controller.traceContextKeys: + - from: sub + to: user.id + - thread_id + template: controller-deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: KAGENT_TRACE_CONTEXT_KEYS + value: '[{"from":"sub","to":"user.id"},"thread_id"]' diff --git a/helm/kagent/values.yaml b/helm/kagent/values.yaml index 7d4b2adedc..0043a92cc2 100644 --- a/helm/kagent/values.yaml +++ b/helm/kagent/values.yaml @@ -206,6 +206,10 @@ config: # ============================================================================== controller: + # Allowlisted caller context copied onto agent spans. Entries may be strings + # or {from, to} mappings, e.g. [{from: sub, to: user.id}, thread_id]. + # Empty disables promotion. Use opaque identifiers; never names or emails. + traceContextKeys: [] # Standard OTEL overrides and external SDK credentials for this process only. env: [] envFrom: [] diff --git a/python/packages/kagent-core/pyproject.toml b/python/packages/kagent-core/pyproject.toml index af9b797055..e4624805c4 100644 --- a/python/packages/kagent-core/pyproject.toml +++ b/python/packages/kagent-core/pyproject.toml @@ -13,6 +13,7 @@ dependencies = [ "kagent-proto>=0.1.0", "opentelemetry-api>=1.42.1,<2", "opentelemetry-sdk>=1.42.1,<2", + "opentelemetry-processor-baggage>=0.63b0,<1", "opentelemetry-exporter-otlp-proto-grpc>=1.42.1,<2", "opentelemetry-exporter-otlp-proto-http>=1.42.1,<2", "opentelemetry-distro>=0.63b0,<1", diff --git a/python/packages/kagent-core/src/kagent/core/a2a/_telemetry.py b/python/packages/kagent-core/src/kagent/core/a2a/_telemetry.py index 88a6856b2b..7df61d0e76 100644 --- a/python/packages/kagent-core/src/kagent/core/a2a/_telemetry.py +++ b/python/packages/kagent-core/src/kagent/core/a2a/_telemetry.py @@ -17,6 +17,7 @@ from opentelemetry.trace import SpanKind, Status, StatusCode from ..telemetry import _boot +from ..telemetry._context_attributes import context_with_promoted_metadata from ..telemetry._conv import ( A2A_METHOD, A2A_TASK_ID, @@ -146,7 +147,8 @@ def _state_name(event: Any) -> str: return a2a_pb2.TaskState.Name(state) if state is not None else "" -def _start(method: str, call_context: ServerCallContext) -> Invocation: +def _start(method: str, call_context: ServerCallContext, message: a2a_pb2.Message | None = None) -> Invocation: + promoted = context_with_promoted_metadata(message=message) attributes: dict[str, str] = {A2A_METHOD: method} name = TRANSPORT_SPAN_NAME duration_attributes = None @@ -167,10 +169,15 @@ def _start(method: str, call_context: ServerCallContext) -> Invocation: kind=SpanKind.INTERNAL, attributes=attributes, start_time=start_ns, + context=promoted, record_exception=False, set_status_on_exception=False, ) + # Processors run after initial span attributes. Restore runtime identity + # in case an allowlisted baggage destination names the same attribute. + span.set_attributes(attributes) invocation = Invocation(span, start_ns, duration_attributes) + invocation.context = trace.set_span_in_context(span, promoted) call_context.state[_INVOCATION_STATE] = invocation return invocation @@ -218,7 +225,12 @@ async def execute(self, context: RequestContext, event_queue) -> None: state = context.call_context.state if context.call_context else {} invocation = state.get(_INVOCATION_STATE) if not isinstance(invocation, Invocation) or invocation.finished: - await self.executor.execute(context, event_queue) + promoted = context_with_promoted_metadata(message=context.message) + token = otel_context.attach(promoted) + try: + await self.executor.execute(context, event_queue) + finally: + otel_context.detach(token) return invocation.adopted = True identity = request_attributes(context) @@ -245,7 +257,7 @@ class TelemetryRequestHandlerMixin: """Open the invocation span around message sends and close it when execution never adopts it.""" async def on_message_send(self, params, context: ServerCallContext): - invocation = _start(A2aMethodValues.SEND_MESSAGE.value, context) + invocation = _start(A2aMethodValues.SEND_MESSAGE.value, context, params.message) token = otel_context.attach(invocation.context) try: result = await super().on_message_send(params, context) @@ -261,7 +273,7 @@ async def on_message_send(self, params, context: ServerCallContext): return result async def on_message_send_stream(self, params, context: ServerCallContext): - invocation = _start(A2aMethodValues.SEND_STREAMING_MESSAGE.value, context) + invocation = _start(A2aMethodValues.SEND_STREAMING_MESSAGE.value, context, params.message) events = super().on_message_send_stream(params, context) try: token = otel_context.attach(invocation.context) diff --git a/python/packages/kagent-core/src/kagent/core/telemetry/_boot.py b/python/packages/kagent-core/src/kagent/core/telemetry/_boot.py index a338778c9c..8fc4c9b938 100644 --- a/python/packages/kagent-core/src/kagent/core/telemetry/_boot.py +++ b/python/packages/kagent-core/src/kagent/core/telemetry/_boot.py @@ -8,6 +8,7 @@ from opentelemetry import _logs, metrics, trace from opentelemetry.distro import OpenTelemetryConfigurator from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor +from opentelemetry.processor.baggage import BaggageSpanProcessor from opentelemetry.propagate import set_global_textmap from opentelemetry.propagators.composite import CompositePropagator from opentelemetry.sdk.resources import OTELResourceDetector @@ -17,6 +18,7 @@ from .._runtime_config import RuntimeConfig, TelemetryConfig from . import _defaults +from ._context_attributes import allowed_baggage_key_predicate from ._conv import GEN_AI_AGENT_ID, GEN_AI_AGENT_NAME, KAGENT_RUNTIME, KagentRuntimeValues from ._profiles import COMMON_INSTRUMENTORS, Profile, profile_for from ._span_processor import KagentAttributesSpanProcessor @@ -152,10 +154,15 @@ def _resource_attributes(identity: Identity, config: TelemetryConfig) -> dict[st def _install_providers(identity: Identity, config: TelemetryConfig) -> None: if os.getenv("OTEL_SDK_DISABLED", "").strip().lower() == "true": return + processors = [] + predicate = allowed_baggage_key_predicate() + if predicate is not None: + processors.append(BaggageSpanProcessor(predicate)) + processors.append(KagentAttributesSpanProcessor()) try: OpenTelemetryConfigurator().configure( resource_attributes=_resource_attributes(identity, config), - span_processors=[KagentAttributesSpanProcessor()], + span_processors=processors, exporter_args_map=_exporter_args(), ) except Exception: diff --git a/python/packages/kagent-core/src/kagent/core/telemetry/_context_attributes.py b/python/packages/kagent-core/src/kagent/core/telemetry/_context_attributes.py new file mode 100644 index 0000000000..141a01419e --- /dev/null +++ b/python/packages/kagent-core/src/kagent/core/telemetry/_context_attributes.py @@ -0,0 +1,264 @@ +"""Promote allowlisted caller context into W3C baggage.""" + +from __future__ import annotations + +import functools +import json +import os +from contextvars import Token +from dataclasses import dataclass +from typing import Any, Callable, Optional + +from a2a.types import Message +from google.protobuf.json_format import MessageToDict +from opentelemetry import baggage +from opentelemetry import context as otel_context + +# 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 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" + +MAX_CONTEXT_KEYS = 32 +MAX_CONTEXT_KEY_LENGTH = 64 + + +@dataclass(frozen=True) +class _ContextMapping: + source: str + attribute: str + + +def context_with_promoted_metadata( + metadata: Optional[dict[str, Any]] = None, + context: Optional[otel_context.Context] = None, + message: Optional[Message] = None, +) -> 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 ctx + if metadata is None and message is not None: + metadata = MessageToDict(message.metadata) if message.HasField("metadata") else {} + if metadata is None: + metadata = {} + + bag = baggage.get_all(ctx) + for mapping in mappings: + existing = bag.get(mapping.attribute) + if isinstance(existing, str) and existing.strip(): + continue + 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 + ctx = baggage.set_baggage(mapping.attribute, value, context=ctx) + bag = baggage.get_all(ctx) + return ctx + + +def promote_message_metadata_to_baggage( + metadata: Optional[dict[str, Any]] = None, + context: Optional[otel_context.Context] = None, + message: Optional[Message] = None, +) -> Optional[Token[otel_context.Context]]: + """Attach a context with allowlisted metadata written into baggage. + + Returns a token that must be detached with + :func:`detach_promoted_metadata`, or ``None`` when nothing changed. + """ + 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 +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``. + Cached because the allowlist cannot change without a process restart. + """ + raw = os.getenv(TRACE_CONTEXT_KEYS_ENV_VAR, "").strip() + if not raw: + return () + if raw.startswith("["): + return tuple(_cap_mappings(_parse_json_allowlist(raw))) + return tuple(_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): + # 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: + 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: + mappings.append(mapping) + 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.""" + 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) -> Optional[_ContextMapping]: + from_key = from_key.strip() + to_key = to_key.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 + return _ContextMapping(source=from_key, attribute=to_key) + + +def _cap_mappings(mappings: list[_ContextMapping]) -> list[_ContextMapping]: + out: list[_ContextMapping] = [] + seen: set[tuple[str, str]] = set() + for mapping in mappings: + identity = (mapping.source, mapping.attribute) + if identity in seen: + continue + seen.add(identity) + out.append(mapping) + if len(out) == MAX_CONTEXT_KEYS: + break + return out + + +def _is_attribute_key(key: str) -> bool: + 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.""" + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, int): + return str(value) + if isinstance(value, float): + 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.""" + 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: + """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() diff --git a/python/packages/kagent-core/tests/conftest.py b/python/packages/kagent-core/tests/conftest.py new file mode 100644 index 0000000000..bb23206952 --- /dev/null +++ b/python/packages/kagent-core/tests/conftest.py @@ -0,0 +1,11 @@ +import pytest + +from kagent.core.telemetry._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_a2a_telemetry.py b/python/packages/kagent-core/tests/test_a2a_telemetry.py index 630a2dab1d..036ed4aab0 100644 --- a/python/packages/kagent-core/tests/test_a2a_telemetry.py +++ b/python/packages/kagent-core/tests/test_a2a_telemetry.py @@ -330,3 +330,68 @@ async def test_without_bootstrap_the_span_still_carries_request_identity(monkeyp assert span.name == "a2a.request" assert span.attributes["enduser.id"] == "alice" assert "flush" not in log + + +@pytest.mark.parametrize( + "runtime", + [ + KagentRuntimeValues.ADK_PYTHON, + KagentRuntimeValues.LANGGRAPH, + KagentRuntimeValues.OPENAI_AGENTS, + KagentRuntimeValues.CREWAI, + ], +) +@pytest.mark.parametrize( + "allowlist", ["", '[{"from":"sub","to":"user.id"},"gen_ai.agent.name","gen_ai.conversation.id","enduser.id"]'] +) +async def test_metadata_policy_covers_invocation_and_model(monkeypatch, log, runtime, allowlist): + from google.protobuf.json_format import ParseDict + from opentelemetry import baggage + from opentelemetry import context as otel_context + from opentelemetry.processor.baggage import BaggageSpanProcessor + + from kagent.core.telemetry._context_attributes import _allowed_context_mappings, allowed_baggage_key_predicate + + monkeypatch.setenv("KAGENT_TRACE_CONTEXT_KEYS", allowlist) + _allowed_context_mappings.cache_clear() + exporter = InMemorySpanExporter() + provider = TracerProvider() + predicate = allowed_baggage_key_predicate() + if predicate is not None: + provider.add_span_processor(BaggageSpanProcessor(predicate)) + provider.add_span_processor(KagentAttributesSpanProcessor()) + provider.add_span_processor(SimpleSpanProcessor(exporter)) + monkeypatch.setattr(_telemetry, "_tracer", lambda: provider.get_tracer("kagent.core.a2a")) + use(monkeypatch, runtime) + handler = handler_for(Agent(provider.get_tracer("test"), log)) + request = send("metadata-policy") + ParseDict( + { + "sub": "opaque-subject", + "private_token": "must-not-be-exported", + "gen_ai.agent.name": "forged-agent", + "gen_ai.conversation.id": "forged-context", + "enduser.id": "forged-user", + }, + request.message.metadata, + ) + original = otel_context.get_current() + try: + await stream(handler, request, log) + assert otel_context.get_current() is original + assert baggage.get_baggage("user.id") is None + spans = exporter.get_finished_spans() + assert len(spans) == 2 + for span in spans: + attrs = dict(span.attributes or {}) + assert attrs.get("user.id") == ("opaque-subject" if allowlist else None) + if span.instrumentation_scope.name == "kagent.core.a2a": + assert attrs["gen_ai.agent.name"] == "researcher" + assert attrs["enduser.id"] == "alice" + assert attrs["gen_ai.conversation.id"] != "forged-context" + assert not any(key.startswith("a2a.message.metadata.") for key in attrs) + assert "private_token" not in attrs + assert "must-not-be-exported" not in attrs.values() + finally: + await handler.aclose() + provider.shutdown() 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..5c57531d6d --- /dev/null +++ b/python/packages/kagent-core/tests/test_caller_context_attributes.py @@ -0,0 +1,304 @@ +from unittest.mock import MagicMock + +import pytest +from a2a.types import Message, Role +from google.protobuf.json_format import ParseDict +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 +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from kagent.core.telemetry._context_attributes import ( + MAX_CONTEXT_KEYS, + TRACE_CONTEXT_KEYS_ENV_VAR, + _allowed_context_mappings, + allowed_baggage_key_predicate, + context_with_promoted_metadata, + detach_promoted_metadata, + promote_message_metadata_to_baggage, +) +from kagent.core.telemetry._span_processor import ( + KAGENT_ATTRIBUTES_KEY, + KagentAttributesSpanProcessor, +) + + +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(): + context = baggage.set_baggage(key, value, context) + return context + + +def baggage_values(context: otel_context.Context) -> dict[str, str]: + return {key: value for key, value in baggage.get_all(context).items()} + + +@pytest.fixture(autouse=True) +def _clear_allowlist_cache(): + _allowed_context_mappings.cache_clear() + yield + _allowed_context_mappings.cache_clear() + + +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_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_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"), + (3, "3"), + (1.5, "1.5"), + ], + ) + 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): + set_allowlist(monkeypatch, "note") + assert baggage_values(context_with_promoted_metadata({"note": "line\nbreak\tand\x00nul"})) == { + "note": "linebreakandnul" + } + + 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)) + 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): + set_allowlist(monkeypatch, "") + decode = MagicMock(return_value={"thread_id": "T1"}) + monkeypatch.setattr("kagent.core.telemetry._context_attributes.MessageToDict", 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() + monkeypatch.setattr( + "kagent.core.telemetry._context_attributes.MessageToDict", + lambda _message: {"thread_id": "T1"}, + ) + assert baggage_values(context_with_promoted_metadata(message=message)) == {"thread_id": "T1"} + + +def test_allowed_baggage_key_predicate(monkeypatch): + set_allowlist(monkeypatch, "") + assert allowed_baggage_key_predicate() is None + + 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_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 _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") + + +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()) + + token = promote_message_metadata_to_baggage( + metadata=None, + context=baggage_context({"user.id": "opaque-subject"}), + ) + attrs_token = otel_context.attach(otel_context.set_value(KAGENT_ATTRIBUTES_KEY, {"kagent.user_id": "runtime-user"})) + try: + tracer = provider.get_tracer("test") + with tracer.start_as_current_span("invocation"): + with tracer.start_as_current_span("tool"): + with tracer.start_as_current_span("generate_content"): + pass + finally: + otel_context.detach(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 = otel_context.attach( + otel_context.set_value(KAGENT_ATTRIBUTES_KEY, {"gen_ai.conversation.id": "from-runtime"}) + ) + try: + tracer = provider.get_tracer("test") + with tracer.start_as_current_span("generate_content"): + pass + finally: + otel_context.detach(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" + + +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", + } + + +@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() diff --git a/python/uv.lock b/python/uv.lock index 260eab5e10..33f8211c88 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -1997,6 +1997,7 @@ dependencies = [ { name = "opentelemetry-instrumentation-fastapi" }, { name = "opentelemetry-instrumentation-grpc" }, { name = "opentelemetry-instrumentation-httpx" }, + { name = "opentelemetry-processor-baggage" }, { name = "opentelemetry-sdk" }, { name = "protovalidate" }, { name = "typing-extensions" }, @@ -2031,6 +2032,7 @@ requires-dist = [ { name = "opentelemetry-instrumentation-genai-openai-agents", marker = "extra == 'openai-agents'", specifier = ">=1.2b0,<1.3" }, { name = "opentelemetry-instrumentation-grpc", specifier = ">=0.63b0,<1" }, { name = "opentelemetry-instrumentation-httpx", specifier = ">=0.63b0,<1" }, + { name = "opentelemetry-processor-baggage", specifier = ">=0.63b0,<1" }, { name = "opentelemetry-sdk", specifier = ">=1.42.1,<2" }, { name = "protovalidate", specifier = ">=1.1.0,<2" }, { name = "typing-extensions", specifier = ">=4.16.0" }, @@ -3247,6 +3249,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4f/c6/e5682b1bfb320b32505e88255c34ae1e99fe9fb220cc465c244e91967eac/opentelemetry_instrumentation_httpx-0.66b1-py3-none-any.whl", hash = "sha256:0342a4002c6dbc6c4bf22cc7e698f50f5c8b77f63325c6f40c94ab87e016bf4d", size = 17723, upload-time = "2026-10-06T17:35:36.501Z" }, ] +[[package]] +name = "opentelemetry-processor-baggage" +version = "0.66b1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-sdk" }, + { name = "wrapt" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/7d/e8/6b11138881f0b1870a16891311e34d2a74fb21c8a7637e71cb61757b810d/opentelemetry_processor_baggage-0.66b1.tar.gz", hash = "sha256:200c16de37b9afae685bdc368a261d1c091ac7ef4fb892ef93129e2be5036642", size = 8794, upload-time = "2026-10-06T17:36:44.807Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b2/92/02b52c5ac72ffa4a54c496c83b99c6941090a4fd419b4666a73651222b8a/opentelemetry_processor_baggage-0.66b1-py3-none-any.whl", hash = "sha256:b43acdcb387d8934e04e8f265165b1d632ca62509c693941b31df69bf62d1fe5", size = 9467, upload-time = "2026-10-06T17:36:04.101Z" }, +] + [[package]] name = "opentelemetry-proto" version = "1.45.1"