311 lines
11 KiB
Go
311 lines
11 KiB
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(©).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(©).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)
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|