diff --git a/internal/pipeline/finalize_stage.go b/internal/pipeline/finalize_stage.go index 1355496b..6221535d 100644 --- a/internal/pipeline/finalize_stage.go +++ b/internal/pipeline/finalize_stage.go @@ -2,6 +2,9 @@ package pipeline import ( "context" + "crypto/sha256" + "encoding/hex" + "io" "log/slog" "os" "path/filepath" @@ -66,15 +69,13 @@ func (s *FinalizeStage) Execute(ctx context.Context, state *RunState) error { state.Tool.MediaResults = append(state.Tool.MediaResults, MediaResult{Path: mf.Path, ContentType: ct}) } - // 3. Deduplicate + populate media sizes - s.processMedia(state) - - // 3b. Persist assistant-generated images (Codex image_generation_call) to disk - // BEFORE building the assistant message so MediaRefs are included in the session store. - // Source is state.Observe.AssistantImages, which ObserveStage accumulates across - // every iteration — required because LastResponse holds only the final iteration's - // response (an image emitted mid-loop alongside a tool call would otherwise be lost). - var assistantImageRefs []providers.MediaRef + // 3. Persist assistant-generated images (Codex image_generation_call) to disk + // and fold them into MediaResults BEFORE deduplication, so a generated image + // and an explicitly attached copy of it collapse into one attachment. + // Source is state.Observe.AssistantImages, which ObserveStage accumulates + // across every iteration — required because LastResponse holds only the final + // iteration's response (an image emitted mid-loop alongside a tool call would + // otherwise be lost). if s.deps.PersistAssistantImages != nil && len(state.Observe.AssistantImages) > 0 { workspace := "" if state.Workspace != nil { @@ -85,11 +86,21 @@ func (s *FinalizeStage) Execute(ctx context.Context, state *RunState) error { // the scratch message — we harvest MediaRefs from there. scratch := &providers.Message{Images: state.Observe.AssistantImages} s.deps.PersistAssistantImages(scratch, workspace) - assistantImageRefs = scratch.MediaRefs + for _, ref := range scratch.MediaRefs { + state.Tool.MediaResults = append(state.Tool.MediaResults, MediaResult{ + Path: ref.Path, + ContentType: ref.MimeType, + Prompt: ref.Prompt, + }) + } state.Observe.AssistantImages = nil // prevent double-processing on retries } - // 3c. Build final assistant message with MediaRefs for session persistence. + // 4. Populate sizes and drop duplicates across every producer. + s.processMedia(state) + + // 5. Build the final assistant message from the deduplicated media set, so + // session history and outbound delivery carry exactly the same attachments. assistantMsg := providers.Message{ Role: "assistant", Content: state.Observe.FinalContent, @@ -113,21 +124,8 @@ func (s *FinalizeStage) Execute(ctx context.Context, state *RunState) error { Prompt: mr.Prompt, }) } - // Append persisted assistant image refs (Codex image_generation_call output). - assistantMsg.MediaRefs = append(assistantMsg.MediaRefs, assistantImageRefs...) state.Messages.AppendPending(assistantMsg) - // Surface generated images (Codex image_generation_call) on MediaResults - // too, so they reach RunResult.Media / outbound channel delivery — not just - // session history (MediaRefs were already appended to the message above). - for _, ref := range assistantImageRefs { - mr := MediaResult{Path: ref.Path, ContentType: ref.MimeType, Prompt: ref.Prompt} - if info, statErr := os.Stat(ref.Path); statErr == nil { - mr.Size = info.Size() - } - state.Tool.MediaResults = append(state.Tool.MediaResults, mr) - } - // 4. Flush remaining pending messages to session store. // Capture the pre-flush history length so metadata msgCount reflects // history + newly-persisted pending (matches upstream calibration). @@ -215,14 +213,43 @@ func (s *FinalizeStage) processMedia(state *RunState) { } } - // Deduplicate by path - seen := make(map[string]bool, len(media)) + // Deduplicate. Path equality alone is not enough: an agent that generates an + // image and then attaches a copy of it under a different name produces two + // entries for one picture, and both were delivered — the user received the + // same image twice. Identical bytes collapse regardless of path. + seenPath := make(map[string]bool, len(media)) + seenDigest := make(map[string]bool, len(media)) deduped := make([]MediaResult, 0, len(media)) for _, m := range media { - if !seen[m.Path] { - seen[m.Path] = true - deduped = append(deduped, m) + if m.Path != "" { + if seenPath[m.Path] { + continue + } + seenPath[m.Path] = true + + // Remote or already-cleaned paths simply skip content comparison. + if digest, err := fileDigest(m.Path); err == nil { + if seenDigest[digest] { + continue + } + seenDigest[digest] = true + } } + deduped = append(deduped, m) } state.Tool.MediaResults = deduped } + +// fileDigest returns the SHA-256 of a local file. +func fileDigest(path string) (string, error) { + f, err := os.Open(path) + if err != nil { + return "", err + } + defer f.Close() + h := sha256.New() + if _, err := io.Copy(h, f); err != nil { + return "", err + } + return hex.EncodeToString(h.Sum(nil)), nil +} diff --git a/internal/pipeline/finalize_stage_media_dedupe_test.go b/internal/pipeline/finalize_stage_media_dedupe_test.go new file mode 100644 index 00000000..9097f450 --- /dev/null +++ b/internal/pipeline/finalize_stage_media_dedupe_test.go @@ -0,0 +1,182 @@ +package pipeline + +import ( + "context" + "os" + "path/filepath" + "testing" + + "github.com/nextlevelbuilder/goclaw/internal/providers" +) + +// assistantMediaRefs returns the MediaRefs on the assistant message the stage +// persisted. Execute flushes and clears the buffer, so the flushed batch is the +// only place the final message can be observed. +func assistantMediaRefs(t *testing.T, flushed []providers.Message) []providers.MediaRef { + t.Helper() + for i := len(flushed) - 1; i >= 0; i-- { + if flushed[i].Role == "assistant" { + return flushed[i].MediaRefs + } + } + t.Fatal("no assistant message persisted") + return nil +} + +// An agent that generates an image and then attaches a copy of it under a +// different name produces two entries for one picture. Deduplicating by path +// alone kept both, and the user received the same image twice in the chat. +func TestFinalizeStage_IdenticalImageUnderTwoPathsSendsOnce(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + bytes := []byte("identical-image-bytes") + generated := filepath.Join(dir, "poster_20260815-122109_291214.png") + attached := filepath.Join(dir, "poster.jpg") + for _, p := range []string{generated, attached} { + if err := os.WriteFile(p, bytes, 0644); err != nil { + t.Fatal(err) + } + } + + var flushed []providers.Message + deps := &PipelineDeps{ + FlushMessages: func(_ context.Context, _ string, msgs []providers.Message) error { + flushed = append(flushed, msgs...) + return nil + }, + UpdateMetadata: func(context.Context, string, providers.Usage, providers.Usage, int) error { return nil }, + PersistAssistantImages: func(msg *providers.Message, _ string) { + msg.MediaRefs = append(msg.MediaRefs, providers.MediaRef{ + ID: filepath.Base(generated), + MimeType: "image/png", + Kind: "image", + Path: generated, + }) + msg.Images = nil + }, + } + + stage := NewFinalizeStage(deps) + state := defaultState() + state.Observe.FinalContent = "Đây anh." + // The agent attached its own copy through a tool. + state.Tool.MediaResults = []MediaResult{{Path: attached, ContentType: "image/jpeg"}} + state.Observe.AssistantImages = []providers.ImageContent{ + {MimeType: "image/png", Data: "aWRlbnRpY2FsLWltYWdlLWJ5dGVz"}, + } + + if err := stage.Execute(context.Background(), state); err != nil { + t.Fatalf("Execute() error: %v", err) + } + + if len(state.Tool.MediaResults) != 1 { + paths := make([]string, 0, len(state.Tool.MediaResults)) + for _, m := range state.Tool.MediaResults { + paths = append(paths, m.Path) + } + t.Fatalf("MediaResults = %d entries %v, want 1 (same image delivered twice)", len(paths), paths) + } + // Session history must match what was delivered. + if refs := assistantMediaRefs(t, flushed); len(refs) != 1 { + t.Fatalf("assistant MediaRefs = %d, want 1", len(refs)) + } +} + +// Different pictures must still both be delivered. +func TestFinalizeStage_DistinctImagesAreBothKept(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + generated := filepath.Join(dir, "generated.png") + other := filepath.Join(dir, "other.png") + if err := os.WriteFile(generated, []byte("first-image"), 0644); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(other, []byte("second-image-different"), 0644); err != nil { + t.Fatal(err) + } + + var flushed []providers.Message + deps := &PipelineDeps{ + FlushMessages: func(_ context.Context, _ string, msgs []providers.Message) error { + flushed = append(flushed, msgs...) + return nil + }, + UpdateMetadata: func(context.Context, string, providers.Usage, providers.Usage, int) error { return nil }, + PersistAssistantImages: func(msg *providers.Message, _ string) { + msg.MediaRefs = append(msg.MediaRefs, providers.MediaRef{ + ID: filepath.Base(generated), MimeType: "image/png", Kind: "image", Path: generated, + }) + msg.Images = nil + }, + } + + stage := NewFinalizeStage(deps) + state := defaultState() + state.Tool.MediaResults = []MediaResult{{Path: other, ContentType: "image/png"}} + state.Observe.AssistantImages = []providers.ImageContent{ + {MimeType: "image/png", Data: "Zmlyc3QtaW1hZ2U="}, + } + + if err := stage.Execute(context.Background(), state); err != nil { + t.Fatalf("Execute() error: %v", err) + } + if len(state.Tool.MediaResults) != 2 { + t.Fatalf("MediaResults = %d entries, want 2 distinct images", len(state.Tool.MediaResults)) + } + if refs := assistantMediaRefs(t, flushed); len(refs) != 2 { + t.Fatalf("assistant MediaRefs = %d, want 2", len(refs)) + } +} + +// A path repeated verbatim collapses without needing to read the file. +func TestFinalizeStage_RepeatedPathSendsOnce(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + p := filepath.Join(dir, "doc.pdf") + if err := os.WriteFile(p, []byte("pdf"), 0644); err != nil { + t.Fatal(err) + } + + stage := NewFinalizeStage(&PipelineDeps{ + FlushMessages: func(context.Context, string, []providers.Message) error { return nil }, + UpdateMetadata: func(context.Context, string, providers.Usage, providers.Usage, int) error { return nil }, + }) + state := defaultState() + state.Tool.MediaResults = []MediaResult{ + {Path: p, ContentType: "application/pdf"}, + {Path: p, ContentType: "application/pdf"}, + } + + if err := stage.Execute(context.Background(), state); err != nil { + t.Fatalf("Execute() error: %v", err) + } + if len(state.Tool.MediaResults) != 1 { + t.Fatalf("MediaResults = %d entries, want 1", len(state.Tool.MediaResults)) + } +} + +// Media the runtime cannot read on disk (remote or already cleaned up) must not +// be dropped just because content comparison is impossible. +func TestFinalizeStage_UnreadableMediaIsKept(t *testing.T) { + t.Parallel() + + stage := NewFinalizeStage(&PipelineDeps{ + FlushMessages: func(context.Context, string, []providers.Message) error { return nil }, + UpdateMetadata: func(context.Context, string, providers.Usage, providers.Usage, int) error { return nil }, + }) + state := defaultState() + state.Tool.MediaResults = []MediaResult{ + {Path: "/nonexistent/a.png", ContentType: "image/png"}, + {Path: "/nonexistent/b.png", ContentType: "image/png"}, + } + + if err := stage.Execute(context.Background(), state); err != nil { + t.Fatalf("Execute() error: %v", err) + } + if len(state.Tool.MediaResults) != 2 { + t.Fatalf("MediaResults = %d entries, want 2 (unreadable media must survive)", len(state.Tool.MediaResults)) + } +}