From d468db61d45fcb73967bacf2fdc2a94b0b4f8911 Mon Sep 17 00:00:00 2001 From: jack Date: Thu, 30 Jul 2026 08:56:39 +0800 Subject: [PATCH] fix(telemetry): capture Langfuse trace input and output --- internal/runner/runner.go | 7 +- internal/telemetry/langfuse.go | 50 ++++++++++- .../telemetry/langfuse_multimodal_test.go | 82 ++++++++++++++++++- 3 files changed, 134 insertions(+), 5 deletions(-) diff --git a/internal/runner/runner.go b/internal/runner/runner.go index 03d36fe5..44e2267a 100644 --- a/internal/runner/runner.go +++ b/internal/runner/runner.go @@ -36,9 +36,12 @@ func Run( goalStore *tools.GoalStore, tracer *telemetry.LangfuseTracer, tokenUsage *internalmodel.TokenUsage, -) string { +) (response string) { if tracer != nil { - ctx = tracer.WithNewTrace(ctx, "coding_agent") + ctx = tracer.WithNewTrace(ctx, "coding_agent", messages) + defer func() { + tracer.EndTrace(ctx, response) + }() } if rec != nil { ctx = internalagent.WithToolObservationSink(ctx, func(observation internalagent.ToolObservation) { diff --git a/internal/telemetry/langfuse.go b/internal/telemetry/langfuse.go index ad89b646..c6086a39 100644 --- a/internal/telemetry/langfuse.go +++ b/internal/telemetry/langfuse.go @@ -58,10 +58,13 @@ func NewLangfuseTracer(cfg *config.LangfuseConfig) *LangfuseTracer { } // WithNewTrace creates a new Langfuse trace and returns a context carrying its ID. -func (t *LangfuseTracer) WithNewTrace(ctx context.Context, name string) context.Context { +// The trace records the latest user message so the top-level Langfuse input +// represents the user turn rather than only its child generation inputs. +func (t *LangfuseTracer) WithNewTrace(ctx context.Context, name string, messages []*schema.Message) context.Context { traceID, err := t.client.CreateTrace(&langfuseacl.TraceEventBody{ BaseEventBody: langfuseacl.BaseEventBody{Name: name}, TimeStamp: time.Now(), + Input: traceInput(messages), }) if err != nil { config.Logger().Printf("[langfuse] CreateTrace error: %v\n", err) @@ -70,6 +73,51 @@ func (t *LangfuseTracer) WithNewTrace(ctx context.Context, name string) context. return context.WithValue(ctx, traceIDKey, traceID) } +// EndTrace records the final response on the current Langfuse trace. +func (t *LangfuseTracer) EndTrace(ctx context.Context, output string) { + traceID, _ := ctx.Value(traceIDKey).(string) + if traceID == "" { + return + } + if err := t.client.EndTrace(&langfuseacl.TraceEventBody{ + BaseEventBody: langfuseacl.BaseEventBody{ + ID: traceID, + }, + TimeStamp: time.Now(), + Output: output, + }); err != nil { + config.Logger().Printf("[langfuse] EndTrace error: %v\n", err) + } +} + +// traceInput returns the latest user message in a form safe to send to +// Langfuse. Multimodal payloads retain their text but replace images with a +// marker, preventing image bytes or URLs from becoming trace input. +func traceInput(messages []*schema.Message) string { + for i := len(messages) - 1; i >= 0; i-- { + msg := messages[i] + if msg == nil || msg.Role != schema.User { + continue + } + if len(msg.UserInputMultiContent) == 0 { + return msg.Content + } + + parts := make([]string, 0, len(msg.UserInputMultiContent)) + for _, part := range msg.UserInputMultiContent { + if part.Type == schema.ChatMessagePartTypeText { + if part.Text != "" { + parts = append(parts, part.Text) + } + continue + } + parts = append(parts, telemetryImagePlaceholder) + } + return strings.Join(parts, "\n") + } + return "" +} + // Flush ensures all buffered events are sent to Langfuse. // It blocks at most defaultFlushTimeout to avoid stalling program exit. func (t *LangfuseTracer) Flush() { diff --git a/internal/telemetry/langfuse_multimodal_test.go b/internal/telemetry/langfuse_multimodal_test.go index a10a447d..1d7abd19 100644 --- a/internal/telemetry/langfuse_multimodal_test.go +++ b/internal/telemetry/langfuse_multimodal_test.go @@ -12,14 +12,20 @@ import ( ) type captureLangfuse struct { + trace *langfuseacl.TraceEventBody + endedTrace *langfuseacl.TraceEventBody generation *langfuseacl.GenerationEventBody } -func (c *captureLangfuse) CreateTrace(*langfuseacl.TraceEventBody) (string, error) { +func (c *captureLangfuse) CreateTrace(body *langfuseacl.TraceEventBody) (string, error) { + c.trace = body return "trace", nil } -func (c *captureLangfuse) EndTrace(*langfuseacl.TraceEventBody) error { return nil } +func (c *captureLangfuse) EndTrace(body *langfuseacl.TraceEventBody) error { + c.endedTrace = body + return nil +} func (c *captureLangfuse) CreateSpan(*langfuseacl.SpanEventBody) (string, error) { return "span", nil @@ -40,6 +46,78 @@ func (c *captureLangfuse) CreateEvent(*langfuseacl.EventEventBody) (string, erro func (c *captureLangfuse) Flush() {} +func TestWithNewTraceCapturesLatestPlainUserInput(t *testing.T) { + client := &captureLangfuse{} + tracer := &LangfuseTracer{client: client} + + tracer.WithNewTrace(context.Background(), "coding_agent", []*schema.Message{ + schema.UserMessage("previous request"), + {Role: schema.Assistant, Content: "previous response"}, + schema.UserMessage("latest request"), + }) + + if client.trace == nil { + t.Fatal("trace was not created") + } + if got, want := client.trace.Input, "latest request"; got != want { + t.Fatalf("trace input=%q, want %q", got, want) + } +} + +func TestWithNewTraceCapturesSafeLatestUserInput(t *testing.T) { + base64Secret := "base64-secret-pixels" + urlSecret := "https://private.invalid/screenshot.png?token=secret" + imagePrompt := &schema.Message{ + Role: schema.User, + UserInputMultiContent: []schema.MessageInputPart{ + {Type: schema.ChatMessagePartTypeText, Text: "Describe this screenshot"}, + { + Type: schema.ChatMessagePartTypeImageURL, + Image: &schema.MessageInputImage{MessagePartCommon: schema.MessagePartCommon{ + MIMEType: "image/png", + Base64Data: &base64Secret, + }}, + }, + { + Type: schema.ChatMessagePartTypeImageURL, + Image: &schema.MessageInputImage{MessagePartCommon: schema.MessagePartCommon{ + MIMEType: "image/png", + URL: &urlSecret, + }}, + }, + }, + } + client := &captureLangfuse{} + tracer := &LangfuseTracer{client: client} + + ctx := tracer.WithNewTrace(context.Background(), "coding_agent", []*schema.Message{ + schema.UserMessage("previous request"), + imagePrompt, + }) + if client.trace == nil { + t.Fatal("trace was not created") + } + if got, want := client.trace.Input, "Describe this screenshot\n"+telemetryImagePlaceholder+"\n"+telemetryImagePlaceholder; got != want { + t.Fatalf("trace input=%q, want %q", got, want) + } + if strings.Contains(client.trace.Input, "previous request") { + t.Fatalf("trace input includes an earlier user turn: %q", client.trace.Input) + } + for _, secret := range []string{base64Secret, urlSecret} { + if strings.Contains(client.trace.Input, secret) { + t.Fatalf("trace input leaked image data %q", secret) + } + } + + tracer.EndTrace(ctx, "completed response") + if client.endedTrace == nil { + t.Fatal("trace was not ended") + } + if client.endedTrace.ID != "trace" || client.endedTrace.Output != "completed response" { + t.Fatalf("ended trace=%#v", client.endedTrace) + } +} + func TestBeforeModelRewriteStateRedactsEnhancedToolImagesFromLangfuse(t *testing.T) { base64Secret := "base64-secret-pixels" urlSecret := "https://private.invalid/screenshot.png?token=secret"