diff --git a/docs/verification/kagent-after.dql b/docs/verification/kagent-after.dql new file mode 100644 index 0000000000..37a8fce5bf --- /dev/null +++ b/docs/verification/kagent-after.dql @@ -0,0 +1,30 @@ +// kagent memory observability — AFTER memory.* instrumentation +// Run in the Dynatrace dev tenant against a kagent deployment WITH this change and +// OTEL_TRACING_ENABLED=true. Expected: rows for memory.read / memory.write / +// memory.consolidate, each carrying the memory.* governance attributes. +// +// Diff against kagent-baseline.dql: baseline returns 0 rows, this returns the memory +// span inventory with operation, scope, injection outcome, and SUT descriptor. +fetch spans +| filter isNotNull(span.name) and matchesPhrase(span.name, "memory.") +| filter k8s.namespace.name == "kagent" +| fields timestamp, + span.name, + trace.id, + span.parent_id, + memory.operation, + memory.scope, + memory.source, + memory.index_ref, + memory.injection_result, + memory.item.count, + memory.sut.name, + memory.sut.architecture, + memory.sut.store_backend +| sort timestamp desc +| limit 100 + +// Aggregate view — recall injection outcomes over time (governance signal): +// fetch spans +// | filter span.name == "memory.read" and k8s.namespace.name == "kagent" +// | summarize recalls = count(), by: {memory.injection_result} diff --git a/docs/verification/kagent-baseline.dql b/docs/verification/kagent-baseline.dql new file mode 100644 index 0000000000..e4a19fc268 --- /dev/null +++ b/docs/verification/kagent-baseline.dql @@ -0,0 +1,9 @@ +// kagent memory observability — BASELINE (before memory.* instrumentation) +// Run in the Dynatrace dev tenant against a kagent deployment WITHOUT this change. +// Expected result: 0 rows — kagent emits no dedicated memory spans, only gen_ai.* on +// invoke_agent. This is the gap kagent-dev/kagent#1909 closes. +fetch spans +| filter isNotNull(span.name) and matchesPhrase(span.name, "memory.") +| filter k8s.namespace.name == "kagent" +| summarize span_count = count(), by: {span.name} +| sort span_count desc diff --git a/go/adk/pkg/memory/kagent_service.go b/go/adk/pkg/memory/kagent_service.go index c50d00226a..6f801dafe4 100644 --- a/go/adk/pkg/memory/kagent_service.go +++ b/go/adk/pkg/memory/kagent_service.go @@ -11,7 +11,11 @@ import ( "github.com/go-logr/logr" "github.com/kagent-dev/kagent/go/adk/pkg/embedding" + "github.com/kagent-dev/kagent/go/adk/pkg/telemetry" "github.com/kagent-dev/kagent/go/api/adk" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/codes" + "go.opentelemetry.io/otel/trace" "google.golang.org/adk/v2/memory" adkmodel "google.golang.org/adk/v2/model" adksession "google.golang.org/adk/v2/session" @@ -83,6 +87,10 @@ func New(cfg Config) (*KagentMemoryService, error) { // It extracts content from the session, optionally summarizes it, generates embeddings, // and stores it via the Kagent API. func (s *KagentMemoryService) AddSessionToMemory(ctx context.Context, session adksession.Session) error { + ctx, span := telemetry.StartMemorySpan(ctx, telemetry.SpanMemoryWrite, + telemetry.MemoryOperationSave, telemetry.MemoryScopeUser, s.agentName) + defer span.End() + log := logr.FromContextOrDiscard(ctx) log.V(1).Info("Adding session to memory", "sessionID", session.ID(), "userID", session.UserID()) @@ -90,42 +98,92 @@ func (s *KagentMemoryService) AddSessionToMemory(ctx context.Context, session ad rawContent := s.extractSessionContent(session) if rawContent == "" { log.V(1).Info("No content to add to memory", "sessionID", session.ID()) + span.SetAttributes(attribute.Int(telemetry.AttrMemoryItemCount, 0)) return nil } - // Summarize content if model is available + // Summarize content if model is available. Raw session text is stored as-is + // (source=user); LLM-summarized facts are stored as source=agent_inference. contents := []string{rawContent} + source := telemetry.MemorySourceUser if s.model != nil { summarized, err := s.summarizeContent(ctx, rawContent) if err != nil { log.V(1).Info("Summarization failed, using raw content", "error", err) } else if len(summarized) > 0 { contents = summarized + source = telemetry.MemorySourceAgentInference log.V(1).Info("Summarized content", "items", len(contents)) } } + span.SetAttributes(attribute.String(telemetry.AttrMemorySource, source)) // Generate embeddings embeddings, err := s.embeddingClient.Generate(ctx, contents) if err != nil { - return fmt.Errorf("failed to generate embeddings: %w", err) + return s.recordSpanError(span, fmt.Errorf("failed to generate embeddings: %w", err)) } if len(embeddings) != len(contents) { - return fmt.Errorf("embedding count mismatch: got %d, expected %d", len(embeddings), len(contents)) + return s.recordSpanError(span, fmt.Errorf("embedding count mismatch: got %d, expected %d", len(embeddings), len(contents))) } // Store each content item with its embedding for i, content := range contents { if err := s.storeMemory(ctx, session.UserID(), content, embeddings[i]); err != nil { - return fmt.Errorf("failed to store memory %d: %w", i, err) + return s.recordSpanError(span, fmt.Errorf("failed to store memory %d: %w", i, err)) } } + span.SetAttributes(attribute.Int(telemetry.AttrMemoryItemCount, len(contents))) log.Info("Successfully added session to memory", "sessionID", session.ID(), "items", len(contents)) return nil } +// SaveMemoryItem stores a single user-provided content item as a first-class +// memory write and emits a memory.write span. Unlike AddSessionToMemory it stores +// the content verbatim (source=user) without LLM summarization, so it intentionally +// emits no memory.consolidate span — consolidation is reserved for the summarizing +// session-ingestion path. This is the path the save_memory tool drives, i.e. the +// write an agent actually performs at runtime. +func (s *KagentMemoryService) SaveMemoryItem(ctx context.Context, userID, content string) error { + ctx, span := telemetry.StartMemorySpan(ctx, telemetry.SpanMemoryWrite, + telemetry.MemoryOperationSave, telemetry.MemoryScopeUser, s.agentName) + defer span.End() + span.SetAttributes(attribute.String(telemetry.AttrMemorySource, telemetry.MemorySourceUser)) + + if content == "" { + return s.recordSpanError(span, fmt.Errorf("missing required parameter: content")) + } + + embeddings, err := s.embeddingClient.Generate(ctx, []string{content}) + if err != nil { + return s.recordSpanError(span, fmt.Errorf("failed to generate embedding: %w", err)) + } + var vector []float32 + if len(embeddings) > 0 { + vector = embeddings[0] + } + if vector == nil { + return s.recordSpanError(span, fmt.Errorf("embedding generation returned no vectors")) + } + + if err := s.storeMemory(ctx, userID, content, vector); err != nil { + return s.recordSpanError(span, fmt.Errorf("failed to save memory: %w", err)) + } + + span.SetAttributes(attribute.Int(telemetry.AttrMemoryItemCount, 1)) + return nil +} + +// recordSpanError marks span as failed, records err on it, and returns err so +// callers can `return s.recordSpanError(span, err)` in a single line. +func (s *KagentMemoryService) recordSpanError(span trace.Span, err error) error { + span.RecordError(err) + span.SetStatus(codes.Error, err.Error()) + return err +} + // storeMemory stores a single memory item via the Kagent API. func (s *KagentMemoryService) storeMemory(ctx context.Context, userID, content string, vector []float32) error { req := addSessionRequest{ @@ -164,10 +222,17 @@ func (s *KagentMemoryService) storeMemory(ctx context.Context, userID, content s // SearchMemory implements memory.Service.SearchMemory. // It searches for relevant memories using vector similarity. func (s *KagentMemoryService) SearchMemory(ctx context.Context, req *memory.SearchRequest) (*memory.SearchResponse, error) { + // Recall runs before LLM dispatch, so this span attaches as a child of the + // active invoke_agent span when one is present in ctx. + ctx, span := telemetry.StartMemorySpan(ctx, telemetry.SpanMemoryRead, + telemetry.MemoryOperationPrefetch, telemetry.MemoryScopeUser, s.agentName) + defer span.End() + log := logr.FromContextOrDiscard(ctx) log.V(1).Info("Searching memory", "query", req.Query, "userID", req.UserID) if req.Query == "" { + setMemoryReadResult(span, 0) return &memory.SearchResponse{Memories: []memory.Entry{}}, nil } @@ -176,6 +241,8 @@ func (s *KagentMemoryService) SearchMemory(ctx context.Context, req *memory.Sear embeddings, err := s.embeddingClient.Generate(ctx, []string{req.Query}) if err != nil { log.Error(err, "Failed to generate query embedding, returning empty results") + span.RecordError(err) + setMemoryReadResult(span, 0) return &memory.SearchResponse{Memories: []memory.Entry{}}, nil } var vector []float32 @@ -183,6 +250,7 @@ func (s *KagentMemoryService) SearchMemory(ctx context.Context, req *memory.Sear vector = embeddings[0] } if vector == nil { + setMemoryReadResult(span, 0) return &memory.SearchResponse{Memories: []memory.Entry{}}, nil } @@ -197,31 +265,31 @@ func (s *KagentMemoryService) SearchMemory(ctx context.Context, req *memory.Sear body, err := json.Marshal(searchReq) if err != nil { - return nil, fmt.Errorf("failed to marshal request: %w", err) + return nil, s.recordSpanError(span, fmt.Errorf("failed to marshal request: %w", err)) } // Make HTTP request url := fmt.Sprintf("%s/api/memories/search", s.apiURL) httpReq, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body)) if err != nil { - return nil, fmt.Errorf("failed to create request: %w", err) + return nil, s.recordSpanError(span, fmt.Errorf("failed to create request: %w", err)) } httpReq.Header.Set("Content-Type", "application/json") resp, err := s.client.Do(httpReq) if err != nil { - return nil, fmt.Errorf("failed to make request: %w", err) + return nil, s.recordSpanError(span, fmt.Errorf("failed to make request: %w", err)) } defer resp.Body.Close() if resp.StatusCode >= 400 { - return nil, fmt.Errorf("API returned status %d", resp.StatusCode) + return nil, s.recordSpanError(span, fmt.Errorf("API returned status %d", resp.StatusCode)) } // Parse response var results []searchResultItem if err := json.NewDecoder(resp.Body).Decode(&results); err != nil { - return nil, fmt.Errorf("failed to decode response: %w", err) + return nil, s.recordSpanError(span, fmt.Errorf("failed to decode response: %w", err)) } // Convert to memory.Entry @@ -238,16 +306,36 @@ func (s *KagentMemoryService) SearchMemory(ctx context.Context, req *memory.Sear }) } + setMemoryReadResult(span, len(memories)) log.Info("Found memories", "count", len(memories), "query", req.Query) return &memory.SearchResponse{Memories: memories}, nil } +// setMemoryReadResult stamps the recall outcome onto a memory.read span: the number +// of items returned and whether any memory passed the pgvector min-score gate +// (injected) or all were filtered out (filtered). +func setMemoryReadResult(span trace.Span, count int) { + injection := telemetry.MemoryInjectionFiltered + if count > 0 { + injection = telemetry.MemoryInjectionInjected + } + span.SetAttributes( + attribute.Int(telemetry.AttrMemoryItemCount, count), + attribute.String(telemetry.AttrMemoryInjectionResult, injection), + ) +} + // summarizeContent uses the LLM to extract key facts from conversation content. // Returns a list of summarized facts, or the original content wrapped in a slice if summarization fails. func (s *KagentMemoryService) summarizeContent(ctx context.Context, content string) ([]string, error) { + ctx, span := telemetry.StartMemorySpan(ctx, telemetry.SpanMemoryConsolidate, + telemetry.MemoryOperationExtract, telemetry.MemoryScopeUser, s.agentName) + defer span.End() + log := logr.FromContextOrDiscard(ctx) if content == "" { + span.SetAttributes(attribute.Int(telemetry.AttrMemoryItemCount, 0)) return nil, nil } @@ -289,7 +377,7 @@ Summary (JSON List):`, content) var summaryText strings.Builder for resp, err := range iter { if err != nil { - return nil, fmt.Errorf("failed to generate summary: %w", err) + return nil, s.recordSpanError(span, fmt.Errorf("failed to generate summary: %w", err)) } if resp.Content != nil { @@ -304,6 +392,7 @@ Summary (JSON List):`, content) summary := strings.TrimSpace(summaryText.String()) if summary == "" { log.V(1).Info("Empty summary returned, using original content") + span.SetAttributes(attribute.Int(telemetry.AttrMemoryItemCount, 1)) return []string{content}, nil } @@ -317,16 +406,19 @@ Summary (JSON List):`, content) var facts []string if err := json.Unmarshal([]byte(summary), &facts); err != nil { log.V(1).Info("Failed to parse summary as JSON, using original content", "error", err, "output", summary) + span.SetAttributes(attribute.Int(telemetry.AttrMemoryItemCount, 1)) return []string{content}, nil } // Validate all items are strings if slices.Contains(facts, "") { log.V(1).Info("Summary contains empty strings, using original content") + span.SetAttributes(attribute.Int(telemetry.AttrMemoryItemCount, 1)) return []string{content}, nil } log.V(1).Info("Successfully summarized content", "facts", len(facts)) + span.SetAttributes(attribute.Int(telemetry.AttrMemoryItemCount, len(facts))) return facts, nil } diff --git a/go/adk/pkg/memory/kagent_service_spans_test.go b/go/adk/pkg/memory/kagent_service_spans_test.go new file mode 100644 index 0000000000..42de23fe9b --- /dev/null +++ b/go/adk/pkg/memory/kagent_service_spans_test.go @@ -0,0 +1,202 @@ +package memory + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/kagent-dev/kagent/go/adk/pkg/telemetry" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" + "google.golang.org/adk/v2/memory" + adksession "google.golang.org/adk/v2/session" +) + +// recordSpans installs an in-memory span recorder as the global tracer provider for +// the duration of the test and returns it so the caller can assert on emitted spans. +func recordSpans(t *testing.T) *tracetest.SpanRecorder { + t.Helper() + sr := tracetest.NewSpanRecorder() + prev := otel.GetTracerProvider() + otel.SetTracerProvider(sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(sr))) + t.Cleanup(func() { otel.SetTracerProvider(prev) }) + return sr +} + +func spanByName(spans []sdktrace.ReadOnlySpan, name string) sdktrace.ReadOnlySpan { + for _, s := range spans { + if s.Name() == name { + return s + } + } + return nil +} + +func attrValue(span sdktrace.ReadOnlySpan, key string) (attribute.Value, bool) { + for _, kv := range span.Attributes() { + if string(kv.Key) == key { + return kv.Value, true + } + } + return attribute.Value{}, false +} + +func assertAttr(t *testing.T, span sdktrace.ReadOnlySpan, key, want string) { + t.Helper() + v, ok := attrValue(span, key) + if !ok { + t.Fatalf("span %q missing attribute %q", span.Name(), key) + } + if v.AsString() != want { + t.Errorf("span %q attribute %q = %q, want %q", span.Name(), key, v.AsString(), want) + } +} + +func TestSearchMemory_EmitsReadSpan(t *testing.T) { + tests := []struct { + name string + results []searchResultItem + wantInjection string + wantCount int64 + }{ + { + name: "hits_injected", + results: []searchResultItem{{ID: "m1", Content: "fact", Score: 0.9}}, + wantInjection: telemetry.MemoryInjectionInjected, + wantCount: 1, + }, + { + name: "no_hits_filtered", + results: []searchResultItem{}, + wantInjection: telemetry.MemoryInjectionFiltered, + wantCount: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sr := recordSpans(t) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(tt.results) + })) + defer server.Close() + + embClient, embServer := newMockEmbeddingClient(t) + defer embServer.Close() + + svc := &KagentMemoryService{ + agentName: "test-agent", + apiURL: server.URL, + client: server.Client(), + ttlDays: 15, + embeddingClient: embClient, + } + + if _, err := svc.SearchMemory(context.Background(), &memory.SearchRequest{Query: "q", UserID: "u1"}); err != nil { + t.Fatalf("SearchMemory() error = %v", err) + } + + span := spanByName(sr.Ended(), telemetry.SpanMemoryRead) + if span == nil { + t.Fatalf("expected a %q span", telemetry.SpanMemoryRead) + } + assertAttr(t, span, telemetry.AttrMemoryOperation, telemetry.MemoryOperationPrefetch) + assertAttr(t, span, telemetry.AttrMemoryScope, telemetry.MemoryScopeUser) + assertAttr(t, span, telemetry.AttrMemoryIndexRef, "test-agent") + assertAttr(t, span, telemetry.AttrMemorySUTName, "kagent") + assertAttr(t, span, telemetry.AttrMemorySUTStoreBackend, "pgvector") + assertAttr(t, span, telemetry.AttrMemoryInjectionResult, tt.wantInjection) + + v, ok := attrValue(span, telemetry.AttrMemoryItemCount) + if !ok || v.AsInt64() != tt.wantCount { + t.Errorf("%s = %v (ok=%v), want %d", telemetry.AttrMemoryItemCount, v.AsInt64(), ok, tt.wantCount) + } + }) + } +} + +func TestSaveMemoryItem_EmitsWriteSpan(t *testing.T) { + sr := recordSpans(t) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + embClient, embServer := newMockEmbeddingClient(t) + defer embServer.Close() + + svc := &KagentMemoryService{ + agentName: "test-agent", + apiURL: server.URL, + client: server.Client(), + ttlDays: 15, + embeddingClient: embClient, + } + + if err := svc.SaveMemoryItem(context.Background(), "user1", "the secret code is BLUE-PANGOLIN-42"); err != nil { + t.Fatalf("SaveMemoryItem() error = %v", err) + } + + span := spanByName(sr.Ended(), telemetry.SpanMemoryWrite) + if span == nil { + t.Fatalf("expected a %q span from the save_memory path", telemetry.SpanMemoryWrite) + } + assertAttr(t, span, telemetry.AttrMemoryOperation, telemetry.MemoryOperationSave) + // Explicit saves store content verbatim -> source=user, no summarization. + assertAttr(t, span, telemetry.AttrMemorySource, telemetry.MemorySourceUser) + assertAttr(t, span, telemetry.AttrMemoryScope, telemetry.MemoryScopeUser) + assertAttr(t, span, telemetry.AttrMemoryIndexRef, "test-agent") + assertAttr(t, span, telemetry.AttrMemorySUTStoreBackend, "pgvector") + + if v, ok := attrValue(span, telemetry.AttrMemoryItemCount); !ok || v.AsInt64() != 1 { + t.Errorf("%s = %v (ok=%v), want 1", telemetry.AttrMemoryItemCount, v.AsInt64(), ok) + } + + // The verbatim save path must NOT emit a consolidate span. + if c := spanByName(sr.Ended(), telemetry.SpanMemoryConsolidate); c != nil { + t.Errorf("save_memory path unexpectedly emitted a %q span", telemetry.SpanMemoryConsolidate) + } +} + +func TestAddSessionToMemory_EmitsWriteSpan(t *testing.T) { + sr := recordSpans(t) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + embClient, embServer := newMockEmbeddingClient(t) + defer embServer.Close() + + svc := &KagentMemoryService{ + agentName: "test-agent", + apiURL: server.URL, + client: server.Client(), + ttlDays: 15, + embeddingClient: embClient, + } + + session := newMockSession("sess1", "user1", []*adksession.Event{ + newMockEvent("user", "remember this"), + }) + if err := svc.AddSessionToMemory(context.Background(), session); err != nil { + t.Fatalf("AddSessionToMemory() error = %v", err) + } + + span := spanByName(sr.Ended(), telemetry.SpanMemoryWrite) + if span == nil { + t.Fatalf("expected a %q span", telemetry.SpanMemoryWrite) + } + assertAttr(t, span, telemetry.AttrMemoryOperation, telemetry.MemoryOperationSave) + // No model configured, so raw session content is stored -> source=user. + assertAttr(t, span, telemetry.AttrMemorySource, telemetry.MemorySourceUser) + assertAttr(t, span, telemetry.AttrMemorySUTArchitecture, "vector") +} diff --git a/go/adk/pkg/memory/save_memory_tool.go b/go/adk/pkg/memory/save_memory_tool.go index b0fa27aceb..6025980d8a 100644 --- a/go/adk/pkg/memory/save_memory_tool.go +++ b/go/adk/pkg/memory/save_memory_tool.go @@ -1,8 +1,6 @@ package memory import ( - "fmt" - adkagent "google.golang.org/adk/v2/agent" "google.golang.org/adk/v2/tool" "google.golang.org/adk/v2/tool/functiontool" @@ -18,27 +16,12 @@ func NewSaveMemoryTool(svc *KagentMemoryService) (tool.Tool, error) { Name: "save_memory", Description: "Saves a specific piece of information or text to long-term memory. Use this to remember important facts, user preferences, or specific details for future reference.", }, func(toolCtx adkagent.Context, in saveMemoryInput) (map[string]any, error) { - if in.Content == "" { - return nil, fmt.Errorf("missing required parameter: content") - } - - // Generate embedding for the content. - embeddings, err := svc.embeddingClient.Generate(toolCtx, []string{in.Content}) - if err != nil { - return nil, fmt.Errorf("failed to generate embedding: %w", err) - } - var vector []float32 - if len(embeddings) > 0 { - vector = embeddings[0] + // SaveMemoryItem emits the memory.write span for this explicit save. The tool + // is a thin adapter so the span-bearing write path stays unit-testable without + // constructing a full ADK ToolContext. + if err := svc.SaveMemoryItem(toolCtx, toolCtx.UserID(), in.Content); err != nil { + return nil, err } - if vector == nil { - return nil, fmt.Errorf("embedding generation returned no vectors") - } - - if err := svc.storeMemory(toolCtx, toolCtx.UserID(), in.Content, vector); err != nil { - return nil, fmt.Errorf("failed to save memory: %w", err) - } - return map[string]any{"status": "Successfully saved information to long-term memory."}, nil }) } diff --git a/go/adk/pkg/telemetry/memory.go b/go/adk/pkg/telemetry/memory.go new file mode 100644 index 0000000000..00666b116a --- /dev/null +++ b/go/adk/pkg/telemetry/memory.go @@ -0,0 +1,96 @@ +package telemetry + +import ( + "context" + + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/trace" +) + +// Memory span names. +// +// These spans sit alongside the existing gen_ai.* spans on invoke_agent and give +// operators dedicated visibility into the memory subsystem. Read spans (memory.read) +// are started with the caller's context so they attach as children of the active +// invocation span when recall happens before LLM dispatch, keeping the trace tree +// connected. +const ( + SpanMemoryWrite = "memory.write" + SpanMemoryRead = "memory.read" + SpanMemoryConsolidate = "memory.consolidate" +) + +// MemoryOperation values. Aligned with the memory.* semantic-convention proposal +// in kagent-dev/kagent#1909. kagent emits only the operations it actually performs; +// governance operations such as promote/revoke are reserved by the convention but +// not emitted until the memory subsystem models them. +const ( + MemoryOperationSave = "save" + MemoryOperationLoad = "load" + MemoryOperationPrefetch = "prefetch" + MemoryOperationExtract = "extract" +) + +// MemoryScope values. kagent scopes memory by user within an agent namespace. +const ( + MemoryScopeUser = "user" +) + +// MemorySource values. Describes where a stored memory originated. +const ( + MemorySourceUser = "user" // raw session content + MemorySourceAgentInference = "agent_inference" // LLM-summarized facts +) + +// MemoryInjectionResult values. Reports the outcome of a recall against the +// pgvector min-score gate. +const ( + MemoryInjectionInjected = "injected" // at least one memory passed the threshold and was returned + MemoryInjectionFiltered = "filtered" // no memory passed the threshold +) + +// Memory attribute keys. +// +// The full governance vocabulary (memory.status, memory.authority, memory.record_id +// on write) is defined by the convention but intentionally NOT emitted here: kagent's +// memory subsystem does not yet model lifecycle status or injection authority, so +// emitting them would mean fabricating values. They are reserved for when kagent gains +// a memory governance model. +const ( + AttrMemoryOperation = "memory.operation" + AttrMemoryScope = "memory.scope" + AttrMemorySource = "memory.source" + AttrMemoryIndexRef = "memory.index_ref" + AttrMemoryInjectionResult = "memory.injection_result" + AttrMemoryItemCount = "memory.item.count" + + // SUT (system-under-test) descriptor attributes for the memory backend. + AttrMemorySUTName = "memory.sut.name" + AttrMemorySUTArchitecture = "memory.sut.architecture" + AttrMemorySUTStoreBackend = "memory.sut.store_backend" +) + +const memoryTracerName = "github.com/kagent-dev/kagent/go/adk/pkg/memory" + +// StartMemorySpan starts a memory.* span as a child of any span already present in +// ctx (for example invoke_agent), stamping the operation plus the SUT descriptor +// attributes that identify kagent's pgvector-backed memory subsystem. indexRef +// identifies the logical index the operation targets (the agent memory namespace); +// pass "" to omit. When telemetry is disabled the returned span is a no-op. +func StartMemorySpan(ctx context.Context, spanName, operation, scope, indexRef string) (context.Context, trace.Span) { + ctx, span := otel.Tracer(memoryTracerName).Start(ctx, spanName) + span.SetAttributes( + attribute.String(AttrMemoryOperation, operation), + attribute.String(AttrMemorySUTName, "kagent"), + attribute.String(AttrMemorySUTArchitecture, "vector"), + attribute.String(AttrMemorySUTStoreBackend, "pgvector"), + ) + if scope != "" { + span.SetAttributes(attribute.String(AttrMemoryScope, scope)) + } + if indexRef != "" { + span.SetAttributes(attribute.String(AttrMemoryIndexRef, indexRef)) + } + return ctx, span +} diff --git a/go/core/internal/database/client_postgres.go b/go/core/internal/database/client_postgres.go index 4e0bdc66d9..0d8908bc2b 100644 --- a/go/core/internal/database/client_postgres.go +++ b/go/core/internal/database/client_postgres.go @@ -18,9 +18,37 @@ import ( dbgen "github.com/kagent-dev/kagent/go/core/internal/database/gen" "github.com/kagent-dev/kagent/go/core/pkg/a2acompat/trpcv0" "github.com/pgvector/pgvector-go" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/codes" + semconv "go.opentelemetry.io/otel/semconv/v1.39.0" + "go.opentelemetry.io/otel/trace" "trpc.group/trpc-go/trpc-a2a-go/protocol" ) +// dbTracer emits child spans for the controller's pgvector memory CRUD, so a +// memory trace propagated from the Python/Go ADK runtime continues past the +// otelhttp server span into the actual database work. +// +// Guardrail: never record raw embedding vectors or full query/content text on +// span attributes — cardinality explosion + PII. Counts, scores, limits, the +// collection name, and the operation are the only memory-derived attributes. +var dbTracer = otel.Tracer("kagent.controller.database") + +// Span attribute keys. The db.* attributes reuse the controller's pinned OTel +// semantic conventions (go.opentelemetry.io/otel/semconv/v1.39.0) via the +// semconv constants/helpers below; memory.item.count and memory.limit are +// kagent memory extensions with no semconv equivalent, so they stay raw keys. +const ( + attrMemoryItemCount = "memory.item.count" + attrMemoryLimit = "memory.limit" + + memoryCollection = "memory" + dbOperationSelect = "SELECT" + dbOperationInsert = "INSERT" + dbOperationDelete = "DELETE" +) + type postgresClient struct { q *dbgen.Queries db *pgxpool.Pool @@ -646,6 +674,17 @@ func (c *postgresClient) GetCrewAIFlowState(ctx context.Context, userID, threadI // ── Agent Memory (vector search) ────────────────────────────────────────────── func (c *postgresClient) StoreAgentMemory(ctx context.Context, memory *dbpkg.Memory) error { + ctx, span := dbTracer.Start(ctx, "db.memory.insert", + trace.WithSpanKind(trace.SpanKindClient), + trace.WithAttributes( + semconv.DBSystemNamePostgreSQL, + semconv.DBOperationName(dbOperationInsert), + semconv.DBCollectionName(memoryCollection), + attribute.Int(attrMemoryItemCount, 1), + ), + ) + defer span.End() + id, err := c.q.InsertMemory(ctx, dbgen.InsertMemoryParams{ AgentName: &memory.AgentName, UserID: &memory.UserID, @@ -656,6 +695,8 @@ func (c *postgresClient) StoreAgentMemory(ctx context.Context, memory *dbpkg.Mem AccessCount: &memory.AccessCount, }) if err != nil { + span.RecordError(err) + span.SetStatus(codes.Error, "insert memory") return err } memory.ID = id @@ -663,7 +704,18 @@ func (c *postgresClient) StoreAgentMemory(ctx context.Context, memory *dbpkg.Mem } func (c *postgresClient) StoreAgentMemories(ctx context.Context, memories []*dbpkg.Memory) error { - return c.withTx(ctx, func(q *dbgen.Queries) error { + ctx, span := dbTracer.Start(ctx, "db.memory.insert", + trace.WithSpanKind(trace.SpanKindClient), + trace.WithAttributes( + semconv.DBSystemNamePostgreSQL, + semconv.DBOperationName(dbOperationInsert), + semconv.DBCollectionName(memoryCollection), + attribute.Int(attrMemoryItemCount, len(memories)), + ), + ) + defer span.End() + + err := c.withTx(ctx, func(q *dbgen.Queries) error { for _, m := range memories { id, err := q.InsertMemory(ctx, dbgen.InsertMemoryParams{ AgentName: &m.AgentName, @@ -681,9 +733,25 @@ func (c *postgresClient) StoreAgentMemories(ctx context.Context, memories []*dbp } return nil }) + if err != nil { + span.RecordError(err) + span.SetStatus(codes.Error, "insert memories") + } + return err } func (c *postgresClient) SearchAgentMemory(ctx context.Context, agentName, userID string, embedding pgvector.Vector, limit int) ([]dbpkg.AgentMemorySearchResult, error) { + ctx, span := dbTracer.Start(ctx, "db.memory.search", + trace.WithSpanKind(trace.SpanKindClient), + trace.WithAttributes( + semconv.DBSystemNamePostgreSQL, + semconv.DBOperationName(dbOperationSelect), + semconv.DBCollectionName(memoryCollection), + attribute.Int(attrMemoryLimit, limit), + ), + ) + defer span.End() + rows, err := c.q.SearchAgentMemory(ctx, dbgen.SearchAgentMemoryParams{ Embedding: embedding, AgentName: &agentName, @@ -691,6 +759,8 @@ func (c *postgresClient) SearchAgentMemory(ctx context.Context, agentName, userI Limit: int32(limit), }) if err != nil { + span.RecordError(err) + span.SetStatus(codes.Error, "search agent memory") return nil, fmt.Errorf("failed to search agent memory: %w", err) } @@ -712,6 +782,7 @@ func (c *postgresClient) SearchAgentMemory(ctx context.Context, agentName, userI Score: score, } } + span.SetAttributes(attribute.Int(attrMemoryItemCount, len(results))) // Access-count bookkeeping is best-effort: a failure must not fail the search. if len(results) > 0 { @@ -728,6 +799,16 @@ func (c *postgresClient) SearchAgentMemory(ctx context.Context, agentName, userI } func (c *postgresClient) ListAgentMemories(ctx context.Context, agentName, userID string) ([]dbpkg.Memory, error) { + ctx, span := dbTracer.Start(ctx, "db.memory.list", + trace.WithSpanKind(trace.SpanKindClient), + trace.WithAttributes( + semconv.DBSystemNamePostgreSQL, + semconv.DBOperationName(dbOperationSelect), + semconv.DBCollectionName(memoryCollection), + ), + ) + defer span.End() + normalized := strings.ReplaceAll(agentName, "-", "_") rows, err := c.q.ListAgentMemories(ctx, dbgen.ListAgentMemoriesParams{ AgentName: &agentName, @@ -735,20 +816,35 @@ func (c *postgresClient) ListAgentMemories(ctx context.Context, agentName, userI UserID: &userID, }) if err != nil { + span.RecordError(err) + span.SetStatus(codes.Error, "list agent memories") return nil, fmt.Errorf("failed to list agent memories: %w", err) } memories := make([]dbpkg.Memory, len(rows)) for i, r := range rows { memories[i] = *toMemory(r) } + span.SetAttributes(attribute.Int(attrMemoryItemCount, len(memories))) return memories, nil } func (c *postgresClient) DeleteAgentMemory(ctx context.Context, agentName, userID string) error { + ctx, span := dbTracer.Start(ctx, "db.memory.delete", + trace.WithSpanKind(trace.SpanKindClient), + trace.WithAttributes( + semconv.DBSystemNamePostgreSQL, + semconv.DBOperationName(dbOperationDelete), + semconv.DBCollectionName(memoryCollection), + ), + ) + defer span.End() + if err := c.q.DeleteAgentMemory(ctx, dbgen.DeleteAgentMemoryParams{ AgentName: &agentName, UserID: &userID, }); err != nil { + span.RecordError(err) + span.SetStatus(codes.Error, "delete agent memory") return fmt.Errorf("failed to delete agent memory: %w", err) } normalized := strings.ReplaceAll(agentName, "-", "_") @@ -757,6 +853,8 @@ func (c *postgresClient) DeleteAgentMemory(ctx context.Context, agentName, userI AgentName: &normalized, UserID: &userID, }); err != nil { + span.RecordError(err) + span.SetStatus(codes.Error, "delete normalized agent memory") return fmt.Errorf("failed to delete normalized agent memory: %w", err) } } diff --git a/go/core/internal/database/memory_trace_test.go b/go/core/internal/database/memory_trace_test.go new file mode 100644 index 0000000000..4bdfef30a6 --- /dev/null +++ b/go/core/internal/database/memory_trace_test.go @@ -0,0 +1,118 @@ +package database + +import ( + "context" + "testing" + + dbpkg "github.com/kagent-dev/kagent/go/api/database" + "github.com/pgvector/pgvector-go" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" +) + +// TestSearchAgentMemory_EmitsDBSpan verifies that the controller's pgvector +// search emits a `db.memory.search` child span carrying the expected non-PII +// DB semantic-convention attributes — and, crucially, that it never leaks the +// raw embedding vector or query/content text onto span attributes +// (observability PII guardrail). +func TestSearchAgentMemory_EmitsDBSpan(t *testing.T) { + db := setupTestDB(t) + client := NewClient(db) + + exporter := tracetest.NewInMemoryExporter() + tp := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter)) + prevTracer := dbTracer + dbTracer = tp.Tracer("kagent.controller.database") + t.Cleanup(func() { + dbTracer = prevTracer + _ = tp.Shutdown(context.Background()) + }) + + // Empty table → zero results; the span is emitted regardless of row count. + embedding := pgvector.NewVector(make([]float32, memoryVectorDim)) + _, err := client.SearchAgentMemory(context.Background(), "trace-agent", "trace-user", embedding, 5) + require.NoError(t, err) + + span := findSpan(t, exporter.GetSpans(), "db.memory.search") + + gotAttrs := attrMap(span) + assert.Equal(t, "postgresql", gotAttrs["db.system.name"]) + assert.Equal(t, "SELECT", gotAttrs["db.operation.name"]) + assert.Equal(t, "memory", gotAttrs["db.collection.name"]) + assert.Equal(t, int64(5), gotAttrs["memory.limit"]) + assert.Equal(t, int64(0), gotAttrs["memory.item.count"], "empty table should yield zero rows") + + assertNoPII(t, span) +} + +// TestStoreAgentMemory_EmitsDBSpan verifies the insert path emits a +// `db.memory.insert` span with a per-item count and no PII. +func TestStoreAgentMemory_EmitsDBSpan(t *testing.T) { + db := setupTestDB(t) + client := NewClient(db) + + exporter := tracetest.NewInMemoryExporter() + tp := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter)) + prevTracer := dbTracer + dbTracer = tp.Tracer("kagent.controller.database") + t.Cleanup(func() { + dbTracer = prevTracer + _ = tp.Shutdown(context.Background()) + }) + + err := client.StoreAgentMemory(context.Background(), &dbpkg.Memory{ + AgentName: "trace-agent", + UserID: "trace-user", + Content: "a secret the span must never carry", + Embedding: pgvector.NewVector(make([]float32, memoryVectorDim)), + Metadata: "{}", + }) + require.NoError(t, err) + + span := findSpan(t, exporter.GetSpans(), "db.memory.insert") + + gotAttrs := attrMap(span) + assert.Equal(t, "postgresql", gotAttrs["db.system.name"]) + assert.Equal(t, "INSERT", gotAttrs["db.operation.name"]) + assert.Equal(t, "memory", gotAttrs["db.collection.name"]) + assert.Equal(t, int64(1), gotAttrs["memory.item.count"]) + + assertNoPII(t, span) +} + +// memoryVectorDim is the embedding dimension enforced by the memory schema. +const memoryVectorDim = 768 + +func findSpan(t *testing.T, spans tracetest.SpanStubs, name string) tracetest.SpanStub { + t.Helper() + for _, s := range spans { + if s.Name == name { + return s + } + } + t.Fatalf("no span named %q found (got %d spans)", name, len(spans)) + return tracetest.SpanStub{} +} + +func attrMap(span tracetest.SpanStub) map[string]interface{} { + m := make(map[string]interface{}) + for _, a := range span.Attributes { + m[string(a.Key)] = a.Value.AsInterface() + } + return m +} + +// assertNoPII fails if any span attribute key hints at raw content, query text, +// or the embedding vector — the cardinality/PII footguns the guardrail forbids. +func assertNoPII(t *testing.T, span tracetest.SpanStub) { + t.Helper() + forbidden := []string{"embedding", "vector", "content", "query", "text"} + for _, a := range span.Attributes { + key := string(a.Key) + for _, bad := range forbidden { + assert.NotContains(t, key, bad, "span attribute %q must not expose %q", key, bad) + } + } +} diff --git a/helm/kagent/templates/controller-configmap.yaml b/helm/kagent/templates/controller-configmap.yaml index a3b519b6a3..d2f290201d 100644 --- a/helm/kagent/templates/controller-configmap.yaml +++ b/helm/kagent/templates/controller-configmap.yaml @@ -27,6 +27,17 @@ data: # OpenTelemetry Configuration OTEL_TRACING_ENABLED: {{ .Values.otel.tracing.enabled | quote }} OTEL_LOGGING_ENABLED: {{ .Values.otel.logging.enabled | quote }} + # Forwarded to agent pods (see collectOtelEnvFromProcess). Read by the a2a + # Python SDK to gate its @trace_class auto-instrumentation. Defaults to true + # (values.yaml) for upstream parity; set otel.tracing.a2aSdkInstrumentation to + # false to drop those framework-plumbing spans. NB: no `default` filter — sprig + # treats bool false as empty, so `| default true` would clobber an explicit false. + OTEL_INSTRUMENTATION_A2A_SDK_ENABLED: {{ .Values.otel.tracing.a2aSdkInstrumentation | quote }} + # Forwarded to agent pods. Read by the kagent Python runtime (tracing/_utils.py) + # to gate httpx client-transport spans. Defaults to true (values.yaml) for + # upstream parity; set otel.tracing.httpxClientInstrumentation to false for + # leaner traces. (Same `| quote`-not-`default` rationale as above.) + OTEL_INSTRUMENTATION_HTTPX_CLIENT_ENABLED: {{ .Values.otel.tracing.httpxClientInstrumentation | quote }} {{- $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/values.yaml b/helm/kagent/values.yaml index c4617d0074..b17b3da359 100644 --- a/helm/kagent/values.yaml +++ b/helm/kagent/values.yaml @@ -917,8 +917,32 @@ oauth2-proxy: # ============================================================================== otel: + # When tracing is enabled, the Go and Python ADK runtimes emit dedicated memory.* + # spans alongside the existing gen_ai.* spans on invoke_agent: + # memory.read (operation=prefetch) - recall; child of invoke_agent + # memory.write (operation=save) - session/tool store to memory + # memory.consolidate (operation=extract) - LLM fact extraction before store tracing: enabled: false + # The a2a Python SDK auto-instruments its own internals (event-queue and + # request-handler plumbing) via @trace_class decorators, giving protocol- and + # queue-level detail for cross-agent (A2A) delegation. Enabled by default to + # match upstream behaviour; the a2a SDK reads OTEL_INSTRUMENTATION_A2A_SDK_ENABLED + # (a2a/utils/telemetry.py) and turns its decorators into no-ops when false. + # Operators who want leaner, high-signal-only traces can set this to false to + # drop the a2a framework-plumbing spans (cross-agent trace continuity still + # holds via the _SubagentInterceptor trace-context hook). + a2aSdkInstrumentation: true + # kagent's Python runtime auto-instruments every outbound httpx client call + # (LLM, embedding, controller memory API) via HTTPXClientInstrumentor. These + # client-transport spans carry outbound-call timing and propagate W3C trace + # context, keeping agent->controller and agent->agent hops stitched into one + # trace. Enabled by default to match upstream behaviour; the runtime reads + # OTEL_INSTRUMENTATION_HTTPX_CLIENT_ENABLED (tracing/_utils.py). Operators who + # want leaner traces can set this to false to drop the raw transport spans — + # trace continuity still holds via the inject_trace_context / _SubagentInterceptor + # hooks, which carry the correlation headers without emitting spans. + httpxClientInstrumentation: true exporter: otlp: endpoint: "" diff --git a/python/packages/kagent-adk/pyproject.toml b/python/packages/kagent-adk/pyproject.toml index 8e7def5837..94b9ec1288 100644 --- a/python/packages/kagent-adk/pyproject.toml +++ b/python/packages/kagent-adk/pyproject.toml @@ -23,6 +23,7 @@ dependencies = [ "anthropic[vertex]>=0.116.0", "fastapi>=0.115.1", "google-adk>=1.28.1,<2", + "opentelemetry-api>=1.0.0", # memory.* span instrumentation (parity with Go ADK) "google-genai>=1.21.1", "google-auth>=2.40.2", "httpx>=0.25.0", diff --git a/python/packages/kagent-adk/src/kagent/adk/_a2a.py b/python/packages/kagent-adk/src/kagent/adk/_a2a.py index 83b986d8fc..56093d3b34 100644 --- a/python/packages/kagent-adk/src/kagent/adk/_a2a.py +++ b/python/packages/kagent-adk/src/kagent/adk/_a2a.py @@ -25,6 +25,7 @@ KAgentTaskStore, get_a2a_max_content_length, ) +from kagent.core.tracing import inject_trace_context from ._agent_executor import A2aAgentExecutor, A2aAgentExecutorConfig from ._lifespan import LifespanManager @@ -97,10 +98,18 @@ def build(self, local=False) -> FastAPI: if not local: token_service = KAgentTokenService(self.app_name) + # Re-inject W3C trace context on every outbound controller call. + # httpx client auto-instrumentation is disabled by default, + # which also stopped injecting `traceparent` — so memory.read/write and + # session calls to the controller were orphaned into a separate root + # trace (the same regression was fixed on the A2A hop). This hook + # restores continuity without re-adding httpx transport spans. + event_hooks = token_service.event_hooks() + event_hooks.setdefault("request", []).append(inject_trace_context) http_client = httpx.AsyncClient( # TODO: add user and agent headers base_url=kagent_url_override or self.kagent_url, - event_hooks=token_service.event_hooks(), + event_hooks=event_hooks, ) if session_db_url: session_service = DatabaseSessionService(db_url=session_db_url) diff --git a/python/packages/kagent-adk/src/kagent/adk/_a2a_telemetry.py b/python/packages/kagent-adk/src/kagent/adk/_a2a_telemetry.py new file mode 100644 index 0000000000..38372e74d5 --- /dev/null +++ b/python/packages/kagent-adk/src/kagent/adk/_a2a_telemetry.py @@ -0,0 +1,58 @@ +"""OpenTelemetry attribute helpers for kagent A2A (cross-agent) delegation. + +A remote-agent (A2A) delegation is already wrapped by the ADK ``execute_tool +`` span. Rather than add a new span layer (which would change the +trace structure), these helpers stamp kagent-specific delegation attributes +onto that active span so it reflects the actual cross-agent call: which +sub-agent was invoked, the conversation lineage (context ids used to keep the +sub-agent session continuous), and how the delegated task resolved. +""" + +from typing import Optional + +from opentelemetry import trace + +# A2A delegation attribute keys. +ATTR_A2A_SUBAGENT_NAME = "a2a.subagent.name" +ATTR_A2A_CONTEXT_ID = "a2a.context_id" +ATTR_A2A_PARENT_CONTEXT_ID = "a2a.parent_context_id" +ATTR_A2A_ROOT_CONTEXT_ID = "a2a.root_context_id" +ATTR_A2A_TASK_ID = "a2a.task.id" +ATTR_A2A_TASK_STATE = "a2a.task.state" + + +def annotate_delegation_request( + subagent_name: str, + context_id: str, + lineage_headers: Optional[dict] = None, +) -> None: + """Stamp the outbound delegation shape onto the active span. + + Records which sub-agent is being called and the conversation lineage + (parent/root context ids) that scopes the sub-agent session. No-op when no + span is recording (tracing disabled). + """ + span = trace.get_current_span() + if not span.is_recording(): + return + span.set_attribute(ATTR_A2A_SUBAGENT_NAME, subagent_name) + if context_id: + span.set_attribute(ATTR_A2A_CONTEXT_ID, context_id) + if lineage_headers: + parent = lineage_headers.get("x-kagent-parent-context-id") + root = lineage_headers.get("x-kagent-root-context-id") + if parent: + span.set_attribute(ATTR_A2A_PARENT_CONTEXT_ID, parent) + if root: + span.set_attribute(ATTR_A2A_ROOT_CONTEXT_ID, root) + + +def annotate_delegation_result(task_id: Optional[str], task_state: Optional[str]) -> None: + """Stamp the resolved task id/state of a delegation onto the active span.""" + span = trace.get_current_span() + if not span.is_recording(): + return + if task_id: + span.set_attribute(ATTR_A2A_TASK_ID, task_id) + if task_state: + span.set_attribute(ATTR_A2A_TASK_STATE, task_state) diff --git a/python/packages/kagent-adk/src/kagent/adk/_memory_service.py b/python/packages/kagent-adk/src/kagent/adk/_memory_service.py index 76b8da81ef..c891c91285 100644 --- a/python/packages/kagent-adk/src/kagent/adk/_memory_service.py +++ b/python/packages/kagent-adk/src/kagent/adk/_memory_service.py @@ -13,6 +13,7 @@ from google.adk.sessions import Session from google.genai import types +from kagent.adk import _memory_telemetry as mtel from kagent.adk.models import KAgentEmbedding from kagent.adk.types import EmbeddingConfig @@ -72,63 +73,80 @@ async def _add_session_to_memory_background(self, session: Session, model: Optio session: The session to add to memory model: Optional ADK model object (e.g., OpenAI, KAgentAnthropicLlm) to use for summarization. """ - try: - # Extract content from session events - raw_content = self._extract_session_content(session) - if not raw_content: - logger.debug("No content to add to memory from session %s", session.id) - return - - logger.debug("Adding session %s to memory for user %s", session.id, session.user_id) - - # Summarize content before embedding - # Returns a list of strings (individual facts/memories) - contents = await self._summarize_session_content_async(raw_content, model=model) - - # Filter out empty content items - valid_contents = [c for c in contents if c] - if not valid_contents: - return - - logger.debug("Generating embeddings for %d content items", len(valid_contents)) - - # Batch generate embeddings - if not self._embedding_client: - logger.warning("No embedding client available for session %s", session.id) - return - vectors = await self._embedding_client.generate(valid_contents) - if not vectors: - logger.warning("Failed to generate embeddings for session %s", session.id) - return - - # Prepare batch items - batch_items = [] - - # Iterate over synced content and vectors - for content_item, vector in zip(valid_contents, vectors, strict=True): - if not vector: - continue - - item: Dict[str, Any] = { - "agent_name": self.agent_name, - "user_id": session.user_id, - "content": content_item, - "vector": vector, - } - if self.ttl_days > 0: - item["ttl_days"] = self.ttl_days - batch_items.append(item) - - if not batch_items: - return + # Session ingestion is a memory write. When a summarization model is + # available the stored facts are LLM-derived (source=agent_inference); + # otherwise the raw session text is stored as-is (source=user). The + # summarization step emits a nested memory.consolidate span. + with mtel.start_memory_span( + mtel.SPAN_MEMORY_WRITE, mtel.MEMORY_OPERATION_SAVE, mtel.MEMORY_SCOPE_USER, self.agent_name + ) as span: + source = mtel.MEMORY_SOURCE_AGENT_INFERENCE if model is not None else mtel.MEMORY_SOURCE_USER + span.set_attribute(mtel.ATTR_MEMORY_SOURCE, source) + try: + # Extract content from session events + raw_content = self._extract_session_content(session) + if not raw_content: + logger.debug("No content to add to memory from session %s", session.id) + span.set_attribute(mtel.ATTR_MEMORY_ITEM_COUNT, 0) + return + + logger.debug("Adding session %s to memory for user %s", session.id, session.user_id) + + # Summarize content before embedding + # Returns a list of strings (individual facts/memories) + contents = await self._summarize_session_content_async(raw_content, model=model) + + # Filter out empty content items + valid_contents = [c for c in contents if c] + if not valid_contents: + span.set_attribute(mtel.ATTR_MEMORY_ITEM_COUNT, 0) + return + + logger.debug("Generating embeddings for %d content items", len(valid_contents)) + + # Batch generate embeddings + if not self._embedding_client: + logger.warning("No embedding client available for session %s", session.id) + span.set_attribute(mtel.ATTR_MEMORY_ITEM_COUNT, 0) + return + with mtel.start_embed_span(self.agent_name, len(valid_contents)): + vectors = await self._embedding_client.generate(valid_contents) + if not vectors: + logger.warning("Failed to generate embeddings for session %s", session.id) + span.set_attribute(mtel.ATTR_MEMORY_ITEM_COUNT, 0) + return + + # Prepare batch items + batch_items = [] + + # Iterate over synced content and vectors + for content_item, vector in zip(valid_contents, vectors, strict=True): + if not vector: + continue - response = await self.client.post("/api/memories/sessions/batch", json={"items": batch_items}) - if response.status_code >= 400: - logger.error("Response body: %s", response.text) - response.raise_for_status() - logger.info("Successfully saved %d memory items via batch API", len(batch_items)) - except Exception as e: - logger.error("Failed to save session %s to memory in background: %s", session.id, e) + item: Dict[str, Any] = { + "agent_name": self.agent_name, + "user_id": session.user_id, + "content": content_item, + "vector": vector, + } + if self.ttl_days > 0: + item["ttl_days"] = self.ttl_days + batch_items.append(item) + + if not batch_items: + span.set_attribute(mtel.ATTR_MEMORY_ITEM_COUNT, 0) + return + + response = await self.client.post("/api/memories/sessions/batch", json={"items": batch_items}) + if response.status_code >= 400: + logger.error("Response body: %s", response.text) + response.raise_for_status() + span.set_attribute(mtel.ATTR_MEMORY_ITEM_COUNT, len(batch_items)) + logger.info("Successfully saved %d memory items via batch API", len(batch_items)) + except Exception as e: + mtel.record_span_error(span, e) + logger.error("Failed to save session %s to memory in background: %s", session.id, e) async def add_memory( self, @@ -149,36 +167,50 @@ async def add_memory( if not content: return - logger.debug("Adding specific content to memory for user %s", user_id) + # An explicit add_memory (the save_memory tool path) is a first-class + # memory write. It stores the content verbatim (source=user, no + # summarization), so it emits memory.write but no memory.consolidate — + # consolidation is reserved for the summarizing session-ingestion path. + with mtel.start_memory_span( + mtel.SPAN_MEMORY_WRITE, mtel.MEMORY_OPERATION_SAVE, mtel.MEMORY_SCOPE_USER, self.agent_name + ) as span: + span.set_attribute(mtel.ATTR_MEMORY_SOURCE, mtel.MEMORY_SOURCE_USER) - # Generate embedding - if not self._embedding_client: - logger.warning("No embedding client available") - return - vector = await self._embedding_client.generate(content) - if not vector: - logger.warning("Failed to generate embedding for memory content") - return + logger.debug("Adding specific content to memory for user %s", user_id) + + # Generate embedding + if not self._embedding_client: + logger.warning("No embedding client available") + span.set_attribute(mtel.ATTR_MEMORY_ITEM_COUNT, 0) + return + with mtel.start_embed_span(self.agent_name, 1): + vector = await self._embedding_client.generate(content) + if not vector: + logger.warning("Failed to generate embedding for memory content") + span.set_attribute(mtel.ATTR_MEMORY_ITEM_COUNT, 0) + return - # Send to Kagent API - payload: Dict[str, Any] = { - "agent_name": self.agent_name, - "user_id": user_id, - "content": content, - "vector": vector, - } - if self.ttl_days > 0: - payload["ttl_days"] = self.ttl_days - - try: - response = await self.client.post("/api/memories/sessions", json=payload) - if response.status_code >= 400: - logger.error("Response body: %s", response.text) - response.raise_for_status() - memory_id = response.json().get("id") - logger.info("Successfully saved memory item (id=%s)", memory_id) - except Exception as e: - logger.error("Failed to save memory: %s", e) + # Send to Kagent API + payload: Dict[str, Any] = { + "agent_name": self.agent_name, + "user_id": user_id, + "content": content, + "vector": vector, + } + if self.ttl_days > 0: + payload["ttl_days"] = self.ttl_days + + try: + response = await self.client.post("/api/memories/sessions", json=payload) + if response.status_code >= 400: + logger.error("Response body: %s", response.text) + response.raise_for_status() + memory_id = response.json().get("id") + span.set_attribute(mtel.ATTR_MEMORY_ITEM_COUNT, 1) + logger.info("Successfully saved memory item (id=%s)", memory_id) + except Exception as e: + mtel.record_span_error(span, e) + logger.error("Failed to save memory: %s", e) async def search_memory( self, @@ -197,48 +229,69 @@ async def search_memory( Returns: SearchMemoryResponse containing matching MemoryEntry objects """ - # Generate embedding for the query - if not self._embedding_client: - logger.warning("No embedding client available for search") - return SearchMemoryResponse(memories=[]) - vector = await self._embedding_client.generate(query) - if not vector: - logger.warning("Failed to generate embedding for search query") - return SearchMemoryResponse(memories=[]) - - payload = { - "agent_name": self.agent_name, - "user_id": user_id, - "vector": vector, - "limit": 5, - "min_score": 0.3, - } - - try: - response = await self.client.post("/api/memories/search", json=payload) - if response.status_code >= 400: - logger.error("Response body: %s", response.text) - response.raise_for_status() - results = response.json() - - memories = [] - for item in results: - content = types.Content( - role="user", - parts=[types.Part(text=item.get("content", ""))], - ) - memory_entry = MemoryEntry(id=item.get("id"), content=content) - memories.append(memory_entry) - - if len(memories) == 0: - logger.warning("No memories found for query: %s", query) + # Recall runs before LLM dispatch, so this span attaches as a child of + # the active invoke_agent span when one is present in context. + _search_limit = 5 + _search_min_score = 0.3 + with mtel.start_memory_span( + mtel.SPAN_MEMORY_READ, mtel.MEMORY_OPERATION_PREFETCH, mtel.MEMORY_SCOPE_USER, self.agent_name + ) as span: + # Stamp the pgvector query shape so the read span reflects the actual + # search parameters, not just the outcome. + span.set_attribute(mtel.ATTR_MEMORY_QUERY_TOP_K, _search_limit) + span.set_attribute(mtel.ATTR_MEMORY_QUERY_MIN_SCORE, _search_min_score) + + # Generate embedding for the query + if not self._embedding_client: + logger.warning("No embedding client available for search") + mtel.set_memory_read_result(span, 0) return SearchMemoryResponse(memories=[]) + # Vectorizing the query dominates recall latency; trace it as an + # explicit child so the embed-vs-search split is visible. + with mtel.start_embed_span(self.agent_name, 1): + vector = await self._embedding_client.generate(query) + if not vector: + logger.warning("Failed to generate embedding for search query") + mtel.set_memory_read_result(span, 0) + return SearchMemoryResponse(memories=[]) + + payload = { + "agent_name": self.agent_name, + "user_id": user_id, + "vector": vector, + "limit": _search_limit, + "min_score": _search_min_score, + } + + try: + response = await self.client.post("/api/memories/search", json=payload) + if response.status_code >= 400: + logger.error("Response body: %s", response.text) + response.raise_for_status() + results = response.json() + + memories = [] + for item in results: + content = types.Content( + role="user", + parts=[types.Part(text=item.get("content", ""))], + ) + memory_entry = MemoryEntry(id=item.get("id"), content=content) + memories.append(memory_entry) + + mtel.set_memory_read_result(span, len(memories)) - logger.info("Successfully retrieved memories for query: %s", query) - return SearchMemoryResponse(memories=memories) - except Exception as e: - logger.error("Failed to search memory: %s", e) - return SearchMemoryResponse(memories=[]) + if len(memories) == 0: + logger.warning("No memories found for query: %s", query) + return SearchMemoryResponse(memories=[]) + + logger.info("Successfully retrieved memories for query: %s", query) + return SearchMemoryResponse(memories=memories) + except Exception as e: + mtel.record_span_error(span, e) + mtel.set_memory_read_result(span, 0) + logger.error("Failed to search memory: %s", e) + return SearchMemoryResponse(memories=[]) def _extract_session_content(self, session: Session) -> str: """Extract text content from session events. @@ -337,60 +390,71 @@ async def _summarize_session_content_async( Summary (JSON List):""" - try: - from google.adk.models.llm_request import LlmRequest - from google.genai.types import Content, Part + # A summarization pass with a real model is a memory.consolidate operation: + # it extracts distinct facts from raw session content. Fallbacks (empty / + # unparseable output) report item.count 0 — no facts were consolidated. + with mtel.start_memory_span( + mtel.SPAN_MEMORY_CONSOLIDATE, mtel.MEMORY_OPERATION_EXTRACT, mtel.MEMORY_SCOPE_USER, self.agent_name + ) as span: + try: + from google.adk.models.llm_request import LlmRequest + from google.genai.types import Content, Part + + # Build LLM request using ADK types + llm_request = LlmRequest( + contents=[ + Content( + role="user", + parts=[Part(text=prompt.format(content=content))], + ) + ], + ) - # Build LLM request using ADK types - llm_request = LlmRequest( - contents=[ - Content( - role="user", - parts=[Part(text=prompt.format(content=content))], - ) - ], - ) - - # Call the model directly - logger.debug("Summarizing session content using model %s", model.model) - response_generator = model.generate_content_async(llm_request, stream=False) - - # Consume the async generator (streaming response) - summary_text = "" - async for chunk in response_generator: - if chunk.content and chunk.content.parts: - summary_text += "".join( - part.text for part in chunk.content.parts if hasattr(part, "text") and part.text - ) + # Call the model directly + logger.debug("Summarizing session content using model %s", model.model) + response_generator = model.generate_content_async(llm_request, stream=False) - summary_text = summary_text.strip() + # Consume the async generator (streaming response) + summary_text = "" + async for chunk in response_generator: + if chunk.content and chunk.content.parts: + summary_text += "".join( + part.text for part in chunk.content.parts if hasattr(part, "text") and part.text + ) - if summary_text: - # Clean up potential markdown formatting if model ignores instruction - if summary_text.startswith("```json"): - summary_text = summary_text[7:] - if summary_text.startswith("```"): - summary_text = summary_text[3:] - if summary_text.endswith("```"): - summary_text = summary_text[:-3] summary_text = summary_text.strip() - try: - extracted_list = json.loads(summary_text) - if isinstance(extracted_list, list) and all(isinstance(item, str) for item in extracted_list): - logger.debug("Summarized session content into %d items", len(extracted_list)) - return extracted_list - else: - logger.warning("LLM returned valid JSON but not a list of strings. Falling back to full text.") - except json.JSONDecodeError: - logger.warning( - "Failed to parse LLM output as JSON. Falling back to full text. Output: %s", summary_text - ) - pass - - logger.warning("Empty summary or invalid format returned, using original content") - return [content] - - except Exception as e: - logger.warning("Failed to summarize session content: %s. Using original.", e) - return [content] + if summary_text: + # Clean up potential markdown formatting if model ignores instruction + if summary_text.startswith("```json"): + summary_text = summary_text[7:] + if summary_text.startswith("```"): + summary_text = summary_text[3:] + if summary_text.endswith("```"): + summary_text = summary_text[:-3] + summary_text = summary_text.strip() + + try: + extracted_list = json.loads(summary_text) + if isinstance(extracted_list, list) and all(isinstance(item, str) for item in extracted_list): + logger.debug("Summarized session content into %d items", len(extracted_list)) + span.set_attribute(mtel.ATTR_MEMORY_ITEM_COUNT, len(extracted_list)) + return extracted_list + else: + logger.warning( + "LLM returned valid JSON but not a list of strings. Falling back to full text." + ) + except json.JSONDecodeError: + logger.warning( + "Failed to parse LLM output as JSON. Falling back to full text. Output: %s", summary_text + ) + pass + + logger.warning("Empty summary or invalid format returned, using original content") + span.set_attribute(mtel.ATTR_MEMORY_ITEM_COUNT, 0) + return [content] + + except Exception as e: + mtel.record_span_error(span, e) + logger.warning("Failed to summarize session content: %s. Using original.", e) + return [content] diff --git a/python/packages/kagent-adk/src/kagent/adk/_memory_telemetry.py b/python/packages/kagent-adk/src/kagent/adk/_memory_telemetry.py new file mode 100644 index 0000000000..358337c930 --- /dev/null +++ b/python/packages/kagent-adk/src/kagent/adk/_memory_telemetry.py @@ -0,0 +1,128 @@ +"""OpenTelemetry span helpers for the kagent memory subsystem. + +These spans mirror the Go ADK memory instrumentation (kagent-dev/kagent#1909, +``go/adk/pkg/telemetry/memory.go``) so both runtimes emit the same +``memory.read`` / ``memory.write`` / ``memory.consolidate`` semantic-convention +spans. Read spans are started with the caller's context so they attach as +children of the active invocation span when recall happens before LLM dispatch. + +kagent emits only the operations it actually performs; governance operations +such as promote/revoke are reserved by the convention but not emitted until the +memory subsystem models them. The full governance vocabulary (memory.status, +memory.authority, memory.record_id on write) is likewise reserved but NOT +emitted here — emitting it would mean fabricating values. +""" + +from contextlib import contextmanager +from typing import Iterator, Optional + +from opentelemetry import trace +from opentelemetry.trace import Status, StatusCode + +_TRACER = trace.get_tracer("kagent.adk.memory") + +# Span names. +SPAN_MEMORY_WRITE = "memory.write" +SPAN_MEMORY_READ = "memory.read" +SPAN_MEMORY_CONSOLIDATE = "memory.consolidate" +SPAN_MEMORY_EMBED = "memory.embed" + +# memory.operation values. Aligned with the memory.* semantic-convention proposal. +MEMORY_OPERATION_SAVE = "save" +MEMORY_OPERATION_LOAD = "load" +MEMORY_OPERATION_PREFETCH = "prefetch" +MEMORY_OPERATION_EXTRACT = "extract" +MEMORY_OPERATION_EMBED = "embed" + +# memory.scope values. kagent scopes memory by user within an agent namespace. +MEMORY_SCOPE_USER = "user" + +# memory.source values. Describes where a stored memory originated. +MEMORY_SOURCE_USER = "user" # raw content stored verbatim +MEMORY_SOURCE_AGENT_INFERENCE = "agent_inference" # LLM-summarized facts + +# memory.injection_result values. Reports the outcome of a recall against the +# pgvector min-score gate. +MEMORY_INJECTION_INJECTED = "injected" # at least one memory passed the threshold +MEMORY_INJECTION_FILTERED = "filtered" # no memory passed the threshold + +# Memory attribute keys. +ATTR_MEMORY_OPERATION = "memory.operation" +ATTR_MEMORY_SCOPE = "memory.scope" +ATTR_MEMORY_SOURCE = "memory.source" +ATTR_MEMORY_INDEX_REF = "memory.index_ref" +ATTR_MEMORY_INJECTION_RESULT = "memory.injection_result" +ATTR_MEMORY_ITEM_COUNT = "memory.item.count" + +# Recall query-shape attributes stamped on memory.read so the span reflects the +# actual pgvector search parameters used (not just the outcome). +ATTR_MEMORY_QUERY_TOP_K = "memory.query.top_k" +ATTR_MEMORY_QUERY_MIN_SCORE = "memory.query.min_score" + +# SUT (system-under-test) descriptor attributes for the memory backend. +ATTR_MEMORY_SUT_NAME = "memory.sut.name" +ATTR_MEMORY_SUT_ARCHITECTURE = "memory.sut.architecture" +ATTR_MEMORY_SUT_STORE_BACKEND = "memory.sut.store_backend" + + +@contextmanager +def start_memory_span( + span_name: str, + operation: str, + scope: str, + index_ref: str, +) -> Iterator[trace.Span]: + """Start a memory.* span as a child of any span already active in context. + + Stamps the operation plus the SUT descriptor attributes that identify + kagent's pgvector-backed memory subsystem. ``index_ref`` identifies the + logical index the operation targets (the agent memory namespace); pass "" + to omit. When tracing is disabled the tracer returns a no-op span. + """ + with _TRACER.start_as_current_span(span_name) as span: + span.set_attribute(ATTR_MEMORY_OPERATION, operation) + span.set_attribute(ATTR_MEMORY_SUT_NAME, "kagent") + span.set_attribute(ATTR_MEMORY_SUT_ARCHITECTURE, "vector") + span.set_attribute(ATTR_MEMORY_SUT_STORE_BACKEND, "pgvector") + if scope: + span.set_attribute(ATTR_MEMORY_SCOPE, scope) + if index_ref: + span.set_attribute(ATTR_MEMORY_INDEX_REF, index_ref) + yield span + + +@contextmanager +def start_embed_span(index_ref: str, item_count: int) -> Iterator[trace.Span]: + """Start a memory.embed child span around embedding generation. + + Memory read/write time is dominated by vectorizing the query/content, not by + the pgvector search or store itself. Emitting this as an explicit child of + the active memory.* span makes that embed-vs-store/search split visible in + the trace (previously the read/write looked like one opaque block). Stamps + memory.operation=embed and the item count being embedded. + """ + with _TRACER.start_as_current_span(SPAN_MEMORY_EMBED) as span: + span.set_attribute(ATTR_MEMORY_OPERATION, MEMORY_OPERATION_EMBED) + span.set_attribute(ATTR_MEMORY_SUT_NAME, "kagent") + span.set_attribute(ATTR_MEMORY_SUT_ARCHITECTURE, "vector") + span.set_attribute(ATTR_MEMORY_SUT_STORE_BACKEND, "pgvector") + if index_ref: + span.set_attribute(ATTR_MEMORY_INDEX_REF, index_ref) + span.set_attribute(ATTR_MEMORY_ITEM_COUNT, item_count) + yield span + + +def set_memory_read_result(span: trace.Span, count: int) -> None: + """Stamp the recall outcome onto a memory.read span: the number of items + returned and whether any memory passed the pgvector min-score gate + (injected) or all were filtered out (filtered). + """ + injection = MEMORY_INJECTION_INJECTED if count > 0 else MEMORY_INJECTION_FILTERED + span.set_attribute(ATTR_MEMORY_ITEM_COUNT, count) + span.set_attribute(ATTR_MEMORY_INJECTION_RESULT, injection) + + +def record_span_error(span: trace.Span, err: Exception) -> None: + """Mark a span as failed and record the exception on it.""" + span.record_exception(err) + span.set_status(Status(StatusCode.ERROR, str(err))) diff --git a/python/packages/kagent-adk/src/kagent/adk/_remote_a2a_tool.py b/python/packages/kagent-adk/src/kagent/adk/_remote_a2a_tool.py index 87fd584060..033f45b239 100644 --- a/python/packages/kagent-adk/src/kagent/adk/_remote_a2a_tool.py +++ b/python/packages/kagent-adk/src/kagent/adk/_remote_a2a_tool.py @@ -52,6 +52,9 @@ KAGENT_HITL_DECISION_TYPE_REJECT, extract_hitl_info_from_task, ) +from opentelemetry.propagate import inject + +from . import _a2a_telemetry as a2atel logger = logging.getLogger("kagent_adk." + __name__) @@ -92,6 +95,18 @@ async def intercept(self, method_name, request_payload, http_kwargs, agent_card, headers = dict(http_kwargs.get("headers", {})) headers[_SOURCE_HEADER] = _SOURCE_SUBAGENT + # Inject W3C trace context (traceparent/tracestate) so the remote agent + # CONTINUES this trace instead of starting a new root. This must live + # here — not in httpx/a2a-sdk auto-instrumentation — because both of + # those are disabled by default for span-noise reduction, and + # disabling them silently dropped cross-agent trace + # propagation. `inject()` reads the currently-active span + # context (the parent agent's invoke_agent/execute_tool span) via the + # global textmap propagator and writes only the correlation headers — + # it creates NO span, so trace continuity is restored without + # reintroducing the outbound client-span noise those changes removed. + inject(headers) + if context: if _USER_ID_CONTEXT_KEY in context.state: headers["x-user-id"] = context.state[_USER_ID_CONTEXT_KEY] @@ -328,6 +343,13 @@ async def _handle_first_call(self, args: dict[str, Any], tool_context: ToolConte # to the same user as the parent agent session. call_context = self._build_call_context(tool_context) + # Reflect the delegation on the active ADK execute_tool span: which + # sub-agent, and the conversation lineage that keeps its session + # continuous. + a2atel.annotate_delegation_request( + self.name, self._last_context_id, self._build_lineage_headers(tool_context) + ) + task: Optional[Task] = None try: async for response in client.send_message(request=message, context=call_context): @@ -347,6 +369,9 @@ async def _handle_first_call(self, args: dict[str, Any], tool_context: ToolConte state = task.status.state if task.status else None + # Stamp how the delegated task resolved onto the active span. + a2atel.annotate_delegation_result(task.id, state.value if state else None) + if state == TaskState.input_required: return self._handle_input_required(task, tool_context) diff --git a/python/packages/kagent-adk/tests/unittests/test_a2a_telemetry.py b/python/packages/kagent-adk/tests/unittests/test_a2a_telemetry.py new file mode 100644 index 0000000000..571267b60f --- /dev/null +++ b/python/packages/kagent-adk/tests/unittests/test_a2a_telemetry.py @@ -0,0 +1,42 @@ +"""Tests for kagent A2A delegation span annotations.""" + +from opentelemetry.sdk.trace import TracerProvider + +from kagent.adk import _a2a_telemetry as a2atel + + +def _tracer(): + return TracerProvider().get_tracer("test") + + +def test_annotate_delegation_request_stamps_subagent_and_lineage(): + tracer = _tracer() + with tracer.start_as_current_span("execute_tool bmad_architect") as span: + a2atel.annotate_delegation_request( + "bmad_architect", + "ctx-123", + { + "x-kagent-parent-context-id": "parent-1", + "x-kagent-root-context-id": "root-0", + }, + ) + attrs = span.attributes + assert attrs[a2atel.ATTR_A2A_SUBAGENT_NAME] == "bmad_architect" + assert attrs[a2atel.ATTR_A2A_CONTEXT_ID] == "ctx-123" + assert attrs[a2atel.ATTR_A2A_PARENT_CONTEXT_ID] == "parent-1" + assert attrs[a2atel.ATTR_A2A_ROOT_CONTEXT_ID] == "root-0" + + +def test_annotate_delegation_result_stamps_task_id_and_state(): + tracer = _tracer() + with tracer.start_as_current_span("execute_tool bmad_architect") as span: + a2atel.annotate_delegation_result("task-9", "completed") + attrs = span.attributes + assert attrs[a2atel.ATTR_A2A_TASK_ID] == "task-9" + assert attrs[a2atel.ATTR_A2A_TASK_STATE] == "completed" + + +def test_annotate_is_noop_without_recording_span(): + # No active recording span -> must not raise. + a2atel.annotate_delegation_request("x", "ctx", None) + a2atel.annotate_delegation_result("t", "completed") diff --git a/python/packages/kagent-adk/tests/unittests/test_memory_spans.py b/python/packages/kagent-adk/tests/unittests/test_memory_spans.py new file mode 100644 index 0000000000..d422303346 --- /dev/null +++ b/python/packages/kagent-adk/tests/unittests/test_memory_spans.py @@ -0,0 +1,128 @@ +"""Tests for kagent memory.* OTel span instrumentation (parity with Go ADK).""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest +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.adk import _memory_telemetry as mtel +from kagent.adk._memory_service import KagentMemoryService + + +@pytest.fixture +def exporter(monkeypatch): + """Route memory spans into an in-memory exporter for assertions.""" + exp = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exp)) + # _TRACER is captured at import; patch it to our isolated provider. + monkeypatch.setattr(mtel, "_TRACER", provider.get_tracer("test")) + return exp + + +def _span_by_name(exporter, name): + for s in exporter.get_finished_spans(): + if s.name == name: + return s + return None + + +def _make_service(): + svc = KagentMemoryService(agent_name="test-agent", http_client=MagicMock(), embedding_config=None, ttl_days=15) + svc._embedding_client = MagicMock() + svc._embedding_client.generate = AsyncMock(return_value=[0.1, 0.2, 0.3]) + return svc + + +def _mock_response(json_value): + resp = MagicMock() + resp.status_code = 200 + resp.raise_for_status = MagicMock() + resp.json = MagicMock(return_value=json_value) + return resp + + +@pytest.mark.asyncio +async def test_add_memory_emits_write_span(exporter): + svc = _make_service() + svc.client.post = AsyncMock(return_value=_mock_response({"id": "m1"})) + + await svc.add_memory(app_name="app", user_id="u1", content="the secret code is BLUE-PANGOLIN-42") + + span = _span_by_name(exporter, mtel.SPAN_MEMORY_WRITE) + assert span is not None, "expected a memory.write span from the add_memory (save_memory) path" + attrs = span.attributes + assert attrs[mtel.ATTR_MEMORY_OPERATION] == mtel.MEMORY_OPERATION_SAVE + # Explicit saves store content verbatim -> source=user, no summarization. + assert attrs[mtel.ATTR_MEMORY_SOURCE] == mtel.MEMORY_SOURCE_USER + assert attrs[mtel.ATTR_MEMORY_SCOPE] == mtel.MEMORY_SCOPE_USER + assert attrs[mtel.ATTR_MEMORY_INDEX_REF] == "test-agent" + assert attrs[mtel.ATTR_MEMORY_SUT_STORE_BACKEND] == "pgvector" + assert attrs[mtel.ATTR_MEMORY_ITEM_COUNT] == 1 + + # The verbatim save path must NOT emit a consolidate span. + assert _span_by_name(exporter, mtel.SPAN_MEMORY_CONSOLIDATE) is None + + +@pytest.mark.asyncio +async def test_search_memory_emits_read_span_injected(exporter): + svc = _make_service() + svc.client.post = AsyncMock(return_value=_mock_response([{"id": "m1", "content": "fact"}])) + + resp = await svc.search_memory(app_name="app", user_id="u1", query="q") + assert len(resp.memories) == 1 + + span = _span_by_name(exporter, mtel.SPAN_MEMORY_READ) + assert span is not None, "expected a memory.read span" + attrs = span.attributes + assert attrs[mtel.ATTR_MEMORY_OPERATION] == mtel.MEMORY_OPERATION_PREFETCH + assert attrs[mtel.ATTR_MEMORY_INJECTION_RESULT] == mtel.MEMORY_INJECTION_INJECTED + assert attrs[mtel.ATTR_MEMORY_ITEM_COUNT] == 1 + assert attrs[mtel.ATTR_MEMORY_SUT_NAME] == "kagent" + + +@pytest.mark.asyncio +async def test_search_memory_emits_read_span_filtered(exporter): + svc = _make_service() + svc.client.post = AsyncMock(return_value=_mock_response([])) + + resp = await svc.search_memory(app_name="app", user_id="u1", query="q") + assert len(resp.memories) == 0 + + span = _span_by_name(exporter, mtel.SPAN_MEMORY_READ) + assert span is not None + assert span.attributes[mtel.ATTR_MEMORY_INJECTION_RESULT] == mtel.MEMORY_INJECTION_FILTERED + assert span.attributes[mtel.ATTR_MEMORY_ITEM_COUNT] == 0 + + +@pytest.mark.asyncio +async def test_search_memory_emits_embed_child_and_query_attrs(exporter): + svc = _make_service() + svc.client.post = AsyncMock(return_value=_mock_response([{"id": "m1", "content": "fact"}])) + + await svc.search_memory(app_name="app", user_id="u1", query="q") + + # The query-vectorization step is now an explicit child of memory.read. + embed = _span_by_name(exporter, mtel.SPAN_MEMORY_EMBED) + assert embed is not None, "expected a memory.embed child span from the recall path" + assert embed.attributes[mtel.ATTR_MEMORY_OPERATION] == mtel.MEMORY_OPERATION_EMBED + assert embed.attributes[mtel.ATTR_MEMORY_ITEM_COUNT] == 1 + + # The read span carries the actual pgvector query shape. + read = _span_by_name(exporter, mtel.SPAN_MEMORY_READ) + assert read.attributes[mtel.ATTR_MEMORY_QUERY_TOP_K] == 5 + assert read.attributes[mtel.ATTR_MEMORY_QUERY_MIN_SCORE] == pytest.approx(0.3) + + +@pytest.mark.asyncio +async def test_add_memory_emits_embed_child(exporter): + svc = _make_service() + svc.client.post = AsyncMock(return_value=_mock_response({"id": "m1"})) + + await svc.add_memory(app_name="app", user_id="u1", content="remember this") + + embed = _span_by_name(exporter, mtel.SPAN_MEMORY_EMBED) + assert embed is not None, "expected a memory.embed child span from the save path" + assert embed.attributes[mtel.ATTR_MEMORY_ITEM_COUNT] == 1 diff --git a/python/packages/kagent-adk/tests/unittests/test_remote_a2a_tool.py b/python/packages/kagent-adk/tests/unittests/test_remote_a2a_tool.py index f3af99c779..d8e1869e5c 100644 --- a/python/packages/kagent-adk/tests/unittests/test_remote_a2a_tool.py +++ b/python/packages/kagent-adk/tests/unittests/test_remote_a2a_tool.py @@ -191,6 +191,34 @@ async def test_no_extra_headers_without_state_key(self): ) assert "authorization" not in headers + async def test_injects_traceparent_from_active_span(self): + """W3C trace context is injected so the remote agent continues THIS + trace rather than starting a new root (regression fix). + + Cross-agent propagation must not depend on httpx/a2a-sdk auto- + instrumentation, which are disabled by default. + """ + from opentelemetry.sdk.trace import TracerProvider + + tracer = TracerProvider().get_tracer("test") + interceptor = _SubagentInterceptor() + with tracer.start_as_current_span("parent") as span: + expected_trace_id = format(span.get_span_context().trace_id, "032x") + expected_span_id = format(span.get_span_context().span_id, "016x") + headers = await self._call_intercept(interceptor, state={}) + + # traceparent format: 00--- + assert "traceparent" in headers + parts = headers["traceparent"].split("-") + assert parts[1] == expected_trace_id + assert parts[2] == expected_span_id + + async def test_no_traceparent_without_active_span(self): + """No active/recording span → no traceparent injected (no invalid header).""" + interceptor = _SubagentInterceptor() + headers = await self._call_intercept(interceptor, state={}) + assert "traceparent" not in headers + # --------------------------------------------------------------------------- # First-call tests 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 2d9bb77fc9..0b31ebef5c 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 ._propagation import inject_trace_context from ._utils import configure -__all__ = ["configure"] +__all__ = ["configure", "inject_trace_context"] diff --git a/python/packages/kagent-core/src/kagent/core/tracing/_propagation.py b/python/packages/kagent-core/src/kagent/core/tracing/_propagation.py new file mode 100644 index 0000000000..e4062a3667 --- /dev/null +++ b/python/packages/kagent-core/src/kagent/core/tracing/_propagation.py @@ -0,0 +1,34 @@ +"""Outbound trace-context propagation for httpx clients. + +httpx client auto-instrumentation is enabled by default, but operators can opt +out (``OTEL_INSTRUMENTATION_HTTPX_CLIENT_ENABLED=false`` / +``otel.tracing.httpxClientInstrumentation=false``) to cut span noise. When it is +disabled, nothing injects W3C trace context (``traceparent`` / ``tracestate``) +on outbound agent->controller calls, so the controller's ``/api/memories/*`` +work (and any other call made through the agent's shared httpx client) starts a +*new root trace* instead of nesting under the active ``memory.read`` / +``memory.write`` span. + +This module provides an httpx request event-hook that re-injects the trace +context — correlation headers only, no extra spans — restoring trace continuity +across the memory hop. It is the same remedy applied to A2A sub-agent +calls in ``_remote_a2a_tool.py``, factored out so it can be attached to any +httpx client. The A2A sub-agent hop has the same regression. +""" + +import httpx +from opentelemetry.propagate import inject + + +async def inject_trace_context(request: httpx.Request) -> None: + """httpx request event-hook: inject W3C trace context into outbound headers. + + ``inject()`` reads the currently-active span context and writes the + ``traceparent`` / ``tracestate`` headers via the global textmap propagator. + When no span is active (e.g. tracing disabled) the carrier stays empty and + the request is left untouched, so this hook is always safe to attach. + """ + carrier: dict[str, str] = {} + inject(carrier) + if carrier: + request.headers.update(carrier) diff --git a/python/packages/kagent-core/src/kagent/core/tracing/_utils.py b/python/packages/kagent-core/src/kagent/core/tracing/_utils.py index f56939e986..70d1d94100 100644 --- a/python/packages/kagent-core/src/kagent/core/tracing/_utils.py +++ b/python/packages/kagent-core/src/kagent/core/tracing/_utils.py @@ -160,9 +160,31 @@ def configure(name: str = "kagent", namespace: str = "kagent", fastapi_app: Fast # check endpoint (high-frequency polling requests) and has little # diagnostic value. _excluded_urls = ".*/\\.well-known/agent-card\\.json" - HTTPXClientInstrumentor().instrument(excluded_urls=_excluded_urls) + + # httpx client-transport spans (one per outbound Ollama/LLM/embedding/ + # controller call) capture outbound-call timing AND carry W3C trace + # context on the wire, so they keep agent->controller / agent->agent + # hops stitched into one trace. Enabled by default to match upstream + # behaviour and preserve that detail; operators who want leaner traces + # can opt out via helm otel.tracing.httpxClientInstrumentation -> this + # env var (trace continuity still holds via the inject_trace_context / + # _SubagentInterceptor hooks, which carry the correlation headers + # without emitting spans). + httpx_client_enabled = os.getenv("OTEL_INSTRUMENTATION_HTTPX_CLIENT_ENABLED", "true").lower() == "true" + if httpx_client_enabled: + HTTPXClientInstrumentor().instrument(excluded_urls=_excluded_urls) + else: + logging.info("httpx client instrumentation disabled (OTEL_INSTRUMENTATION_HTTPX_CLIENT_ENABLED=false)") + if fastapi_app: - FastAPIInstrumentor().instrument_app(fastapi_app, excluded_urls=_excluded_urls) + # Instrument the FastAPI server boundary. The agent-card health-check + # endpoint is still excluded (high-frequency polling, no diagnostic + # value); the standard ASGI request spans are kept for parity with + # upstream. + FastAPIInstrumentor().instrument_app( + fastapi_app, + excluded_urls=_excluded_urls, + ) # Configure logging if enabled if logging_enabled: logging.info("Enabling logging for GenAI events") diff --git a/python/packages/kagent-core/tests/test_tracing_configure.py b/python/packages/kagent-core/tests/test_tracing_configure.py index ac99319c6e..2c6f240348 100644 --- a/python/packages/kagent-core/tests/test_tracing_configure.py +++ b/python/packages/kagent-core/tests/test_tracing_configure.py @@ -82,6 +82,94 @@ def fake_instrument_google(): assert instrument_calls["google_instrumented"] is True +def _stub_tracing_side_effects(monkeypatch, instrument_calls): + """Neutralise the OTLP/exporter side effects of configure() so tests can + focus on which auto-instrumentors are (de)activated.""" + + class FakeOpenAIInstrumentor: + def __init__(self, **kwargs): + pass + + def instrument(self, **kwargs): + pass + + class FakeHTTPXInstrumentor: + def instrument(self, **kwargs): + instrument_calls["httpx_instrument_kwargs"] = kwargs + + class FakeFastAPIInstrumentor: + def instrument_app(self, app, **kwargs): + instrument_calls["fastapi_instrument_kwargs"] = kwargs + + monkeypatch.setattr(_utils, "OpenAIInstrumentor", FakeOpenAIInstrumentor) + monkeypatch.setattr(_utils, "HTTPXClientInstrumentor", FakeHTTPXInstrumentor) + monkeypatch.setattr(_utils, "FastAPIInstrumentor", FakeFastAPIInstrumentor) + monkeypatch.setattr(_utils, "_instrument_anthropic", lambda *a, **k: None) + monkeypatch.setattr(_utils, "_instrument_google_generativeai", lambda: None) + monkeypatch.setattr(_utils, "_create_span_exporter", lambda *a, **k: object()) + monkeypatch.setattr(_utils, "BatchSpanProcessor", lambda *a, **k: object()) + monkeypatch.setattr(_utils, "KagentAttributesSpanProcessor", lambda *a, **k: object()) + + class FakeTracerProvider: + def __init__(self, *a, **k): + pass + + def add_span_processor(self, processor): + pass + + # Stub the provider so configure() neither mutates global state nor registers + # a real atexit shutdown hook (the real SDK provider would). + monkeypatch.setattr(_utils, "TracerProvider", FakeTracerProvider) + monkeypatch.setattr(_utils.trace, "get_tracer_provider", lambda: object()) + monkeypatch.setattr(_utils.trace, "set_tracer_provider", lambda provider: None) + + +def test_configure_httpx_client_instrumentation_enabled_by_default(monkeypatch): + """httpx client-transport spans are on by default (upstream parity); they + carry outbound-call timing and on-wire trace context.""" + monkeypatch.setenv("OTEL_TRACING_ENABLED", "true") + monkeypatch.setenv("OTEL_LOGGING_ENABLED", "false") + monkeypatch.delenv("OTEL_INSTRUMENTATION_HTTPX_CLIENT_ENABLED", raising=False) + + instrument_calls = {} + _stub_tracing_side_effects(monkeypatch, instrument_calls) + + _utils.configure(name="test", namespace="test") + + assert "httpx_instrument_kwargs" in instrument_calls + assert "excluded_urls" in instrument_calls["httpx_instrument_kwargs"] + + +def test_configure_httpx_client_instrumentation_opt_out_via_env(monkeypatch): + """Operators can opt out of the raw transport spans for leaner traces.""" + monkeypatch.setenv("OTEL_TRACING_ENABLED", "true") + monkeypatch.setenv("OTEL_LOGGING_ENABLED", "false") + monkeypatch.setenv("OTEL_INSTRUMENTATION_HTTPX_CLIENT_ENABLED", "false") + + instrument_calls = {} + _stub_tracing_side_effects(monkeypatch, instrument_calls) + + _utils.configure(name="test", namespace="test") + + assert "httpx_instrument_kwargs" not in instrument_calls + + +def test_configure_fastapi_keeps_standard_asgi_spans(monkeypatch): + """FastAPI is instrumented with the standard ASGI spans (upstream parity); + only the agent-card health-check endpoint is excluded via excluded_urls.""" + monkeypatch.setenv("OTEL_TRACING_ENABLED", "true") + monkeypatch.setenv("OTEL_LOGGING_ENABLED", "false") + + instrument_calls = {} + _stub_tracing_side_effects(monkeypatch, instrument_calls) + + _utils.configure(name="test", namespace="test", fastapi_app=object()) + + kwargs = instrument_calls["fastapi_instrument_kwargs"] + assert "exclude_spans" not in kwargs + assert "excluded_urls" in kwargs + + def test_otel_sdk_default_propagator_includes_w3c_tracecontext(): """The OTEL SDK must propagate W3C TraceContext by default. diff --git a/python/packages/kagent-core/tests/test_tracing_propagation.py b/python/packages/kagent-core/tests/test_tracing_propagation.py new file mode 100644 index 0000000000..722486737b --- /dev/null +++ b/python/packages/kagent-core/tests/test_tracing_propagation.py @@ -0,0 +1,49 @@ +"""Tests for the outbound trace-context injection httpx hook (memory hop).""" + +import httpx +import pytest +from opentelemetry.sdk.trace import TracerProvider + +from kagent.core.tracing import inject_trace_context + + +def _make_request() -> httpx.Request: + return httpx.Request("POST", "http://kagent-controller/api/memories/search") + + +@pytest.mark.asyncio +async def test_inject_adds_traceparent_when_span_active(): + # A real (non-recording) provider is enough: an active span yields a valid + # SpanContext that the W3C propagator serializes into `traceparent`. + provider = TracerProvider() + tracer = provider.get_tracer("test") + + request = _make_request() + with tracer.start_as_current_span("memory.read"): + await inject_trace_context(request) + + assert "traceparent" in request.headers + # W3C traceparent format: version-traceid-spanid-flags (3 hyphens). + assert request.headers["traceparent"].count("-") == 3 + + +@pytest.mark.asyncio +async def test_inject_is_noop_without_active_span(): + # No active/recording span -> empty carrier -> request headers untouched. + request = _make_request() + await inject_trace_context(request) + assert "traceparent" not in request.headers + + +@pytest.mark.asyncio +async def test_inject_preserves_existing_headers(): + provider = TracerProvider() + tracer = provider.get_tracer("test") + + request = _make_request() + request.headers["Authorization"] = "Bearer abc" + with tracer.start_as_current_span("memory.read"): + await inject_trace_context(request) + + assert request.headers["Authorization"] == "Bearer abc" + assert "traceparent" in request.headers diff --git a/python/uv.lock b/python/uv.lock index 30c3e9f473..d640fbcd9a 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -2312,6 +2312,7 @@ dependencies = [ { name = "numpy", version = "2.3.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "ollama" }, { name = "openai" }, + { name = "opentelemetry-api" }, { name = "protobuf" }, { name = "pydantic" }, { name = "python-multipart" }, @@ -2353,6 +2354,7 @@ requires-dist = [ { name = "numpy", specifier = ">=2.2.6" }, { name = "ollama", specifier = ">=0.3.6" }, { name = "openai", specifier = ">=1.72.0" }, + { name = "opentelemetry-api", specifier = ">=1.0.0" }, { name = "protobuf", specifier = ">=6.33.5" }, { name = "psutil", marker = "extra == 'memory'", specifier = ">=6.1.0" }, { name = "pydantic", specifier = ">=2.5.0" },