fix(agent): redact error event messages

.github/workflows/prc.yml

@@ -3,6 +3,8 @@ on:
   pull_request:
   merge_group:
   push:
+    branches:
+      - main
   workflow_dispatch:
 permissions:
   contents: read

agent/n8n/n8n_agent.go

@@ -24,6 +24,7 @@ import (
 	"trpc.group/trpc-go/trpc-agent-go/agent"
 	"trpc.group/trpc-go/trpc-agent-go/event"
 	"trpc.group/trpc-go/trpc-agent-go/model"
+	"trpc.group/trpc-go/trpc-agent-go/platform"
 	"trpc.group/trpc-go/trpc-agent-go/tool"
 )
 
@@ -220,12 +221,20 @@ func (a *N8nAgent) sendErrorEvent(
 		a.name,
 		event.WithResponse(&model.Response{
 			Error: &model.ResponseError{
-				Message: errorMessage,
+				Message: redactAgentErrorMessage(errorMessage),
 			},
 		}),
 	))
 }
 
+func redactAgentErrorMessage(errorMessage string) string {
+	redactor, err := platform.NewRedactor()
+	if err != nil {
+		return errorMessage
+	}
+	return redactor.Redact(errorMessage)
+}
+
 func (a *N8nAgent) sendFinalStreamingEvent(
 	ctx context.Context,
 	eventChan chan<- *event.Event,

agent/n8n/n8n_agent_test.go

@@ -13,6 +13,7 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"strings"
 	"testing"
 	"time"
 
@@ -222,7 +223,12 @@ func TestN8nAgent_SendErrorEvent(t *testing.T) {
 	}
 
 	eventChan := make(chan *event.Event, 1)
-	a.sendErrorEvent(context.Background(), eventChan, invocation, "test error message")
+	a.sendErrorEvent(
+		context.Background(),
+		eventChan,
+		invocation,
+		"request failed Authorization: ApiKey raw-token\napi_key=sk-1234567890abcdef\nCookie: session=abc; sid=def",
+	)
 	close(eventChan)
 
 	evt := <-eventChan
@@ -235,8 +241,16 @@ func TestN8nAgent_SendErrorEvent(t *testing.T) {
 	if evt.Response.Error == nil {
 		t.Fatal("expected error in response")
 	}
-	if evt.Response.Error.Message != "test error message" {
-		t.Errorf("expected error message 'test error message', got: %s", evt.Response.Error.Message)
+	message := evt.Response.Error.Message
+	for _, secret := range []string{"raw-token", "sk-1234567890abcdef", "session=abc", "sid=def"} {
+		if strings.Contains(message, secret) {
+			t.Errorf("expected error message to redact %q, got: %s", secret, message)
+		}
+	}
+	for _, redacted := range []string{"Authorization: ****", "api_key=****", "Cookie: ****"} {
+		if !strings.Contains(message, redacted) {
+			t.Errorf("expected error message to contain %q, got: %s", redacted, message)
+		}
 	}
 	if evt.Author != "test-agent" {
 		t.Errorf("expected author 'test-agent', got: %s", evt.Author)

agent/weknora/weknora_agent.go

@@ -21,6 +21,7 @@ import (
 	"trpc.group/trpc-go/trpc-agent-go/agent"
 	"trpc.group/trpc-go/trpc-agent-go/event"
 	"trpc.group/trpc-go/trpc-agent-go/model"
+	"trpc.group/trpc-go/trpc-agent-go/platform"
 	"trpc.group/trpc-go/trpc-agent-go/session"
 	"trpc.group/trpc-go/trpc-agent-go/tool"
 )
@@ -67,12 +68,20 @@ func (r *WeKnoraAgent) sendErrorEvent(ctx context.Context, eventChan chan<- *eve
 		r.name,
 		event.WithResponse(&model.Response{
 			Error: &model.ResponseError{
-				Message: errorMessage,
+				Message: redactAgentErrorMessage(errorMessage),
 			},
 		}),
 	))
 }
 
+func redactAgentErrorMessage(errorMessage string) string {
+	redactor, err := platform.NewRedactor()
+	if err != nil {
+		return errorMessage
+	}
+	return redactor.Redact(errorMessage)
+}
+
 // Run implements the Agent interface
 func (r *WeKnoraAgent) Run(ctx context.Context, invocation *agent.Invocation) (<-chan *event.Event, error) {
 	if invocation != nil && invocation.RunOptions.Stream != nil && !*invocation.RunOptions.Stream {

agent/weknora/weknora_agent_test.go

@@ -14,6 +14,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 	"time"
 
@@ -224,7 +225,12 @@ func TestWeKnoraAgent_SendErrorEvent(t *testing.T) {
 	}
 
 	eventChan := make(chan *event.Event, 1)
-	weknoraAgent.sendErrorEvent(context.Background(), eventChan, invocation, "test error message")
+	weknoraAgent.sendErrorEvent(
+		context.Background(),
+		eventChan,
+		invocation,
+		"weknora agent error Authorization: Bearer raw-token\napi_key=sk-1234567890abcdef\nCookie: session=abc; sid=def",
+	)
 	close(eventChan)
 
 	evt := <-eventChan
@@ -237,8 +243,16 @@ func TestWeKnoraAgent_SendErrorEvent(t *testing.T) {
 	if evt.Response.Error == nil {
 		t.Fatal("expected error in response")
 	}
-	if evt.Response.Error.Message != "test error message" {
-		t.Errorf("expected error message 'test error message', got: %s", evt.Response.Error.Message)
+	message := evt.Response.Error.Message
+	for _, secret := range []string{"raw-token", "sk-1234567890abcdef", "session=abc", "sid=def"} {
+		if strings.Contains(message, secret) {
+			t.Errorf("expected error message to redact %q, got: %s", secret, message)
+		}
+	}
+	for _, redacted := range []string{"Authorization: ****", "api_key=****", "Cookie: ****"} {
+		if !strings.Contains(message, redacted) {
+			t.Errorf("expected error message to contain %q, got: %s", redacted, message)
+		}
 	}
 	if evt.Author != "test-agent" {
 		t.Errorf("expected author 'test-agent', got: %s", evt.Author)

internal/flow/processor/content.go

@@ -28,6 +28,8 @@ import (
 	"trpc.group/trpc-go/trpc-agent-go/graph"
 	"trpc.group/trpc-go/trpc-agent-go/internal/fileref"
 	iflow "trpc.group/trpc-go/trpc-agent-go/internal/flow"
+	itelemetry "trpc.group/trpc-go/trpc-agent-go/internal/telemetry"
+	itrace "trpc.group/trpc-go/trpc-agent-go/internal/trace"
 	"trpc.group/trpc-go/trpc-agent-go/internal/util/message"
 	"trpc.group/trpc-go/trpc-agent-go/log"
 	"trpc.group/trpc-go/trpc-agent-go/memory"
@@ -3107,8 +3109,23 @@ func (p *ContentRequestProcessor) getAdaptivePreloadMemoryMessage(
 		Deduplicate:  true,
 		HybridSearch: true,
 	}
-	memories, err := reader.SearchMemories(
-		ctx,
+	searchCtx, span, startedSpan := itrace.StartSpan(ctx, inv, itelemetry.NewMemorySearchSpanName())
+	var memories []*memory.Entry
+	if startedSpan {
+		defer func() {
+			itelemetry.TraceMemorySearch(
+				span,
+				searchOpts.MaxResults,
+				len(memories),
+				searchOpts.HybridSearch,
+				searchOpts.Deduplicate,
+				err,
+			)
+			span.End()
+		}()
+	}
+	memories, err = reader.SearchMemories(
+		searchCtx,
 		userKey,
 		query,
 		memory.WithSearchOptions(searchOpts),

internal/flow/processor/content_memory_test.go

@@ -17,12 +17,16 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/codes"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
 	"trpc.group/trpc-go/trpc-agent-go/agent"
 	"trpc.group/trpc-go/trpc-agent-go/event"
+	itelemetry "trpc.group/trpc-go/trpc-agent-go/internal/telemetry"
 	"trpc.group/trpc-go/trpc-agent-go/memory"
 	"trpc.group/trpc-go/trpc-agent-go/model"
 	"trpc.group/trpc-go/trpc-agent-go/session"
 	"trpc.group/trpc-go/trpc-agent-go/session/inmemory"
+	semconvtrace "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/trace"
 	"trpc.group/trpc-go/trpc-agent-go/tool"
 )
 
@@ -371,6 +375,49 @@ func (m *mockMemoryService) Close() error {
 	return nil
 }
 
+func requireMemorySearchSpan(t *testing.T, recorder interface {
+	Ended() []sdktrace.ReadOnlySpan
+}) sdktrace.ReadOnlySpan {
+	t.Helper()
+	for _, span := range recorder.Ended() {
+		if span.Name() == itelemetry.NewMemorySearchSpanName() {
+			return span
+		}
+	}
+	t.Fatalf("span %q not recorded; ended spans=%v", itelemetry.NewMemorySearchSpanName(), recorder.Ended())
+	return nil
+}
+
+func requireMemorySearchSpanAttribute(t *testing.T, span sdktrace.ReadOnlySpan, key string, want any) {
+	t.Helper()
+	for _, attr := range span.Attributes() {
+		if string(attr.Key) != key {
+			continue
+		}
+		switch v := want.(type) {
+		case string:
+			require.Equal(t, v, attr.Value.AsString())
+		case int64:
+			require.Equal(t, v, attr.Value.AsInt64())
+		case bool:
+			require.Equal(t, v, attr.Value.AsBool())
+		default:
+			t.Fatalf("unsupported expected attribute type %T", want)
+		}
+		return
+	}
+	t.Fatalf("missing attribute %s=%v; attributes=%v", key, want, span.Attributes())
+}
+
+func requireNoMemorySearchSpanAttribute(t *testing.T, span sdktrace.ReadOnlySpan, key string) {
+	t.Helper()
+	for _, attr := range span.Attributes() {
+		if string(attr.Key) == key {
+			t.Fatalf("unexpected attribute %s present; attributes=%v", key, span.Attributes())
+		}
+	}
+}
+
 type mockSearchableSessionService struct {
 	session.Service
 	searchResults []session.EventSearchResult
@@ -666,6 +713,34 @@ func TestGetPreloadMemoryMessage(t *testing.T) {
 		assert.Contains(t, msg.Content, "Relevant memory")
 	})
 
+	t.Run("positive preload records memory search trace contract", func(t *testing.T) {
+		recorder := useSpanRecorder(t)
+		p := NewContentRequestProcessor(WithPreloadMemory(2))
+		mockSvc := &mockMemoryService{
+			memories: []*memory.Entry{
+				newTestMemoryEntry("mem-1", "first"),
+				newTestMemoryEntry("mem-2", "second"),
+				newTestMemoryEntry("mem-3", "third"),
+			},
+			searchResults: []*memory.Entry{
+				newTestMemoryEntry("mem-search", "Relevant memory"),
+			},
+		}
+		inv := newTestInvocation(model.NewUserMessage("find relevant"), mockSvc)
+
+		msg := p.getPreloadMemoryMessage(context.Background(), inv)
+
+		require.NotNil(t, msg)
+		span := requireMemorySearchSpan(t, recorder)
+		requireMemorySearchSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoTraceSpan, itelemetry.OperationMemorySearch)
+		requireMemorySearchSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemorySearchMaxResults, int64(2))
+		requireMemorySearchSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemorySearchResultCount, int64(1))
+		requireMemorySearchSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemorySearchHybrid, true)
+		requireMemorySearchSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemorySearchDeduplicate, true)
+		requireNoMemorySearchSpanAttribute(t, span, "trpc.go.agent.memory.search.query")
+		require.NotEqual(t, codes.Error, span.Status().Code)
+	})
+
 	t.Run("positive preload falls back to recent load when query is empty", func(t *testing.T) {
 		p := NewContentRequestProcessor(WithPreloadMemory(2))
 		mockSvc := &mockMemoryService{
@@ -705,6 +780,28 @@ func TestGetPreloadMemoryMessage(t *testing.T) {
 		assert.NotContains(t, msg.Content, "third")
 	})
 
+	t.Run("positive preload records memory search trace contract on search error", func(t *testing.T) {
+		recorder := useSpanRecorder(t)
+		p := NewContentRequestProcessor(WithPreloadMemory(2))
+		mockSvc := &mockMemoryService{
+			memories: []*memory.Entry{
+				newTestMemoryEntry("mem-1", "first"),
+				newTestMemoryEntry("mem-2", "second"),
+				newTestMemoryEntry("mem-3", "third"),
+			},
+			searchErr: assert.AnError,
+		}
+		inv := newTestInvocation(model.NewUserMessage("hello"), mockSvc)
+
+		msg := p.getPreloadMemoryMessage(context.Background(), inv)
+
+		require.NotNil(t, msg)
+		span := requireMemorySearchSpan(t, recorder)
+		requireMemorySearchSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoTraceSpan, itelemetry.OperationMemorySearch)
+		requireMemorySearchSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemorySearchResultCount, int64(0))
+		require.Equal(t, codes.Error, span.Status().Code)
+	})
+
 	t.Run("positive preload falls back to recent load when search is empty", func(t *testing.T) {
 		p := NewContentRequestProcessor(WithPreloadMemory(2))
 		mockSvc := &mockMemoryService{

internal/flow/processor/functioncall.go

@@ -797,6 +797,7 @@ func (p *FunctionCallResponseProcessor) executeSingleToolCallSequentialResult(
 ) (toolResult, error) {
 	ctx, span, startedSpan := itrace.StartSpan(ctx, invocation, itelemetry.NewExecuteToolSpanName(toolCall.Function.Name))
 	if startedSpan {
+		itelemetry.MarkToolCallSpan(span)
 		defer span.End()
 	}
 	startTime := time.Now()
@@ -1101,6 +1102,7 @@ func (p *FunctionCallResponseProcessor) runParallelToolCall(
 	// Trace the tool execution for observability.
 	ctx, span, startedSpan := itrace.StartSpan(ctx, invocation, itelemetry.NewExecuteToolSpanName(tc.Function.Name))
 	if startedSpan {
+		itelemetry.MarkToolCallSpan(span)
 		defer span.End()
 	}
 	startTime := time.Now()