xinmi/internal/workbench/content_memory_test.go

package workbench

import (
	"context"
	"strings"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"gorm.io/gorm"
)

type feedbackMemoryStub struct {
	interfaces.MemoryService
	available bool
}

func (m feedbackMemoryStub) MemoryAvailable(context.Context) bool { return m.available }

func feedbackHarness(t *testing.T) (*Service, context.Context, *Resource, *Job) {
	t.Helper()
	return feedbackHarnessWithService(t, sqliteWorkbench(t))
}

func feedbackHarnessWithService(t *testing.T, s *Service) (*Service, context.Context, *Resource, *Job) {
	t.Helper()
	s.agents = generationAgentStub{}
	s.memories = feedbackMemoryStub{available: true}
	ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
	r := seedResource(t, s, 1, "alice", "content-tasks")
	r.Data = asJSON(TaskData{Segments: []Segment{{ID: "segment", AgentID: "author", Status: "generated", Title: "开场", Content: "AI 生成的测试正文不是用户事实", MemoryAdvice: "开场先说结论，再给案例"}}})
	if err := s.db.Save(r).Error; err != nil {
		t.Fatal(err)
	}
	until := time.Now().Add(time.Minute)
	j := &Job{ID: "accept-job", TenantID: 1, OwnerID: "alice", Kind: r.Kind, ResourceID: r.ID, Action: "accept", Status: "running", LeaseToken: "lease", LeaseUntil: &until}
	if err := s.db.Create(j).Error; err != nil {
		t.Fatal(err)
	}
	ctx = context.WithValue(ctx, leaseContextKey{}, jobLease{j.ID, j.LeaseToken})
	return s, ctx, r, j
}

func TestContentFeedbackPersistsOnlyWithAcceptedSourceAndIsIdempotent(t *testing.T) {
	s, ctx, r, j := feedbackHarness(t)
	if err := s.executeTask(ctx, j, r, ActionRequest{SegmentID: "segment"}); err != nil {
		t.Fatal(err)
	}
	var count int64
	s.db.Model(&PendingMemory{}).Count(&count)
	if count != 0 {
		t.Fatal("execution leaked memory before content save")
	}
	if err := s.save(ctx, r, 1); err != nil {
		t.Fatal(err)
	}
	if r.pendingMemory != nil {
		t.Fatal("saved feedback retained a stale runtime candidate")
	}
	var row PendingMemory
	if err := s.db.First(&row).Error; err != nil {
		t.Fatal(err)
	}
	if row.Content != "开场先说结论，再给案例" || strings.Contains(row.Content, "AI 生成") || row.Status != "pending" || row.OwnerID != "alice" || row.TenantID != 1 || row.AgentID != "author" || row.SourceRef != contentMemorySource(r.ID, "segment") {
		t.Fatalf("incorrect persisted candidate: %+v", row)
	}
	if err := s.executeTask(ctx, j, r, ActionRequest{SegmentID: "segment"}); err != nil {
		t.Fatal(err)
	}
	if err := s.save(ctx, r, 2); err != nil {
		t.Fatal(err)
	}
	s.db.Model(&PendingMemory{}).Count(&count)
	if count != 1 {
		t.Fatalf("repeated acceptance created %d candidates", count)
	}
	if err := s.db.Model(&PendingMemory{}).Where("id = ?", row.ID).Update("status", "rejected").Error; err != nil {
		t.Fatal(err)
	}
	if err := s.executeTask(ctx, j, r, ActionRequest{SegmentID: "segment"}); err != nil {
		t.Fatal(err)
	}
	if err := s.save(ctx, r, 3); err != nil {
		t.Fatal(err)
	}
	if err := s.db.First(&row, "id = ?", row.ID).Error; err != nil || row.Status != "rejected" {
		t.Fatalf("rejected feedback resurrected: %+v %v", row, err)
	}
}

func TestContentFeedbackWriteFailureRollsBackAcceptanceAndCanRetry(t *testing.T) {
	s, ctx, r, j := feedbackHarness(t)
	if err := s.db.Exec("CREATE TRIGGER reject_feedback BEFORE INSERT ON workbench_pending_memories BEGIN SELECT RAISE(ABORT, 'temporary failure'); END").Error; err != nil {
		t.Fatal(err)
	}
	if err := s.executeTask(ctx, j, r, ActionRequest{SegmentID: "segment"}); err != nil {
		t.Fatal(err)
	}
	if err := s.save(ctx, r, 1); err == nil {
		t.Fatal("failed memory write reported successful acceptance")
	}
	if err := s.db.First(r, "id = ?", r.ID).Error; err != nil {
		t.Fatal(err)
	}
	d, _ := decode[TaskData](r.Data)
	if r.Version != 1 || d.Segments[0].Status != "generated" {
		t.Fatal("failed candidate write leaked adopted content state")
	}
	if err := s.db.Exec("DROP TRIGGER reject_feedback").Error; err != nil {
		t.Fatal(err)
	}
	if err := s.executeTask(ctx, j, r, ActionRequest{SegmentID: "segment"}); err != nil {
		t.Fatal(err)
	}
	if err := s.save(ctx, r, 1); err != nil {
		t.Fatal(err)
	}
	var count int64
	s.db.Model(&PendingMemory{}).Count(&count)
	if count != 1 {
		t.Fatalf("retry lost or duplicated candidate: %d", count)
	}
}

func TestContentFeedbackSourceChangesInvalidateOnlyPendingOwnCandidates(t *testing.T) {
	for _, change := range []string{"skip", "rewrite", "agent", "removed"} {
		t.Run(change, func(t *testing.T) {
			s, ctx, r, j := feedbackHarness(t)
			if err := s.executeTask(ctx, j, r, ActionRequest{SegmentID: "segment"}); err != nil {
				t.Fatal(err)
			}
			if err := s.save(ctx, r, 1); err != nil {
				t.Fatal(err)
			}
			var candidate PendingMemory
			if err := s.db.First(&candidate).Error; err != nil {
				t.Fatal(err)
			}
			for _, id := range []string{"other-user", "other-tenant", "already-accepted"} {
				copy := candidate
				copy.ID = id
				if id == "other-user" {
					copy.OwnerID = "bob"
				}
				if id == "other-tenant" {
					copy.TenantID = 2
				}
				if id == "already-accepted" {
					copy.Status = "accepted"
				}
				if err := s.db.Create(&copy).Error; err != nil {
					t.Fatal(err)
				}
			}
			d, _ := decode[TaskData](r.Data)
			switch change {
			case "skip":
				d.Segments[0].Status = "skipped"
			case "rewrite":
				d.Segments[0].Content = "新正文"
				d.Segments[0].Status = "generated"
			case "agent":
				d.Segments[0].AgentID = "replacement"
			case "removed":
				d.Segments = nil
			}
			// Reusing an in-memory resource must not revive a stale candidate.
			r.Data, r.pendingMemory = asJSON(d), &candidate
			if err := s.save(ctx, r, 2); err != nil {
				t.Fatal(err)
			}
			if err := s.db.First(&candidate, "id = ?", candidate.ID).Error; err != nil || candidate.Status != "invalidated" {
				t.Fatalf("stale feedback remains reviewable: %+v %v", candidate, err)
			}
			var untouched int64
			s.db.Model(&PendingMemory{}).Where("id IN ? AND status IN ?", []string{"other-user", "other-tenant", "already-accepted"}, []string{"pending", "accepted"}).Count(&untouched)
			if untouched != 3 {
				t.Fatal("source invalidation crossed owner/status boundary")
			}
		})
	}
}

func TestDeletingContentTaskInvalidatesPendingFeedbackAtomically(t *testing.T) {
	for _, fail := range []bool{false, true} {
		t.Run(map[bool]string{false: "deleted", true: "write_failure"}[fail], func(t *testing.T) {
			s, ctx, r, j := feedbackHarness(t)
			if err := s.executeTask(ctx, j, r, ActionRequest{SegmentID: "segment"}); err != nil {
				t.Fatal(err)
			}
			if err := s.save(ctx, r, 1); err != nil {
				t.Fatal(err)
			}
			var candidate PendingMemory
			if err := s.db.First(&candidate).Error; err != nil {
				t.Fatal(err)
			}
			for _, id := range []string{"other-owner", "other-task", "accepted"} {
				copy := candidate
				copy.ID = id
				switch id {
				case "other-owner":
					copy.OwnerID = "bob"
				case "other-task":
					copy.SourceRef = contentMemorySource("another-task", "segment")
				case "accepted":
					copy.Status = "accepted"
				}
				if err := s.db.Create(&copy).Error; err != nil {
					t.Fatal(err)
				}
			}
			if fail {
				if err := s.db.Exec("CREATE TRIGGER reject_invalidation BEFORE UPDATE ON workbench_pending_memories BEGIN SELECT RAISE(ABORT, 'temporary failure'); END").Error; err != nil {
					t.Fatal(err)
				}
			}
			err := s.db.Transaction(func(tx *gorm.DB) error { return deleteResource(ctx, tx, r) })
			if (err != nil) != fail {
				t.Fatalf("unexpected deletion error: %v", err)
			}
			var count int64
			if err := s.db.Model(&Resource{}).Where("id = ?", r.ID).Count(&count).Error; err != nil {
				t.Fatal(err)
			}
			if (count == 1) != fail {
				t.Fatal("deletion was not atomic with memory invalidation")
			}
			if err := s.db.First(&candidate, "id = ?", candidate.ID).Error; err != nil {
				t.Fatal(err)
			}
			if (candidate.Status == "pending") != fail {
				t.Fatalf("unexpected candidate status: %s", candidate.Status)
			}
			if err := s.db.Model(&PendingMemory{}).Where("id IN ? AND status IN ?", []string{"other-owner", "other-task", "accepted"}, []string{"pending", "accepted"}).Count(&count).Error; err != nil || count != 3 {
				t.Fatalf("deletion invalidated unrelated memory: %d %v", count, err)
			}
		})
	}
}

func TestReviewContentFeedbackRequiresCurrentAcceptedSource(t *testing.T) {
	for _, mode := range []string{"current", "missing", "skipped", "changed_advice", "changed_agent", "other_owner", "other_tenant"} {
		t.Run(mode, func(t *testing.T) {
			s, _, r, _ := feedbackHarness(t)
			d, _ := decode[TaskData](r.Data)
			d.Segments[0].Status = "accepted"
			row := PendingMemory{TenantID: r.TenantID, OwnerID: r.OwnerID, AgentID: d.Segments[0].AgentID,
				SourceRef: contentMemorySource(r.ID, d.Segments[0].ID), Content: d.Segments[0].MemoryAdvice}
			switch mode {
			case "missing":
				row.SourceRef = contentMemorySource("deleted-task", d.Segments[0].ID)
			case "skipped":
				d.Segments[0].Status = "skipped"
			case "changed_advice":
				d.Segments[0].MemoryAdvice = "另一条建议"
			case "changed_agent":
				d.Segments[0].AgentID = "other-agent"
			case "other_owner":
				row.OwnerID = "bob"
			case "other_tenant":
				row.TenantID = 2
			}
			r.Data = asJSON(d)
			if err := s.db.Save(r).Error; err != nil {
				t.Fatal(err)
			}
			err := s.db.Transaction(func(tx *gorm.DB) error { return lockContentFeedbackSource(tx, &row) })
			if (err == nil) != (mode == "current") {
				t.Fatalf("unexpected source-review result: %v", err)
			}
		})
	}
}

func TestContentFeedbackRespectsScopeAndMemorySwitches(t *testing.T) {
	for _, mode := range []string{"enabled", "draft", "skipped", "no_advice", "owner", "tenant", "workspace_disabled", "agent_disabled", "snapshot_disabled"} {
		t.Run(mode, func(t *testing.T) {
			s, ctx, r, _ := feedbackHarness(t)
			d, _ := decode[TaskData](r.Data)
			segment := d.Segments[0]
			segment.Status = "accepted"
			disabled := false
			snapshots := map[string]AgentSnapshot{}
			switch mode {
			case "draft":
				segment.Status = "generated"
			case "skipped":
				segment.Status = "skipped"
			case "no_advice":
				segment.MemoryAdvice = ""
			case "owner":
				ctx = workbenchContext(1, "bob", types.TenantRoleAdmin)
			case "tenant":
				ctx = workbenchContext(2, "alice", types.TenantRoleAdmin)
			case "workspace_disabled":
				s.memories = feedbackMemoryStub{}
			case "agent_disabled":
				s.agents = generationAgentStub{config: types.CustomAgentConfig{MemoryEnabled: &disabled}}
			case "snapshot_disabled":
				snapshots[segment.AgentID] = AgentSnapshot{MemoryEnabled: &disabled}
			}
			candidate := s.contentFeedbackCandidate(ctx, r, segment, snapshots)
			if (candidate != nil) != (mode == "enabled") {
				t.Fatalf("unexpected candidate for %s: %+v", mode, candidate)
			}
		})
	}
}