WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/workbench/concurrency_test.go

@@ -0,0 +1,498 @@
+package workbench
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"log"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"path/filepath"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/gin-gonic/gin"
+	"github.com/google/uuid"
+	"gorm.io/driver/postgres"
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+	"gorm.io/gorm/logger"
+)
+
+func workbenchContext(tenant uint64, user string, role types.TenantRole) context.Context {
+	ctx := context.WithValue(context.Background(), types.TenantIDContextKey, tenant)
+	ctx = context.WithValue(ctx, types.UserIDContextKey, user)
+	return context.WithValue(ctx, types.TenantRoleContextKey, role)
+}
+
+func sqliteWorkbench(t *testing.T) *Service {
+	t.Helper()
+	db, err := gorm.Open(sqlite.Open(filepath.Join(t.TempDir(), "workbench.db")+"?_busy_timeout=5000&_journal_mode=WAL"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
+	if err != nil {
+		t.Fatal(err)
+	}
+	sqlDB, err := db.DB()
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = sqlDB.Close() })
+	if err = db.AutoMigrate(&Resource{}, &Job{}, &Revision{}, &Event{}, &Consumption{}, &Balance{}, &BalanceHold{}, &Transaction{}, &Usage{}, &UsageQuota{}, &OperationLog{}, &PendingMemory{}); err != nil {
+		t.Fatal(err)
+	}
+	return NewService(db)
+}
+
+func seedResource(t *testing.T, s *Service, tenant uint64, owner, kind string) *Resource {
+	t.Helper()
+	r := &Resource{ID: uuid.NewString(), TenantID: tenant, OwnerID: owner, Kind: kind, Name: "original", Status: "draft", Version: 1, Data: JSON{}}
+	if err := s.db.Create(r).Error; err != nil {
+		t.Fatal(err)
+	}
+	return r
+}
+
+func TestSaveRequiresCurrentLeaseAndScope(t *testing.T) {
+	for _, state := range []string{"valid", "expired", "token", "status", "resource", "owner", "tenant"} {
+		t.Run(state, func(t *testing.T) {
+			s := sqliteWorkbench(t)
+			ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
+			r := seedResource(t, s, 1, "alice", "training-maps")
+			until := time.Now().Add(time.Minute)
+			j := Job{ID: uuid.NewString(), TenantID: 1, OwnerID: "alice", ResourceID: r.ID, Kind: r.Kind, Status: "running", LeaseToken: "lease", LeaseUntil: &until}
+			switch state {
+			case "expired":
+				until = time.Now().Add(-time.Second)
+			case "token":
+				j.LeaseToken = "other"
+			case "status":
+				j.Status = "failed"
+			case "resource":
+				j.ResourceID = uuid.NewString()
+			case "owner":
+				j.OwnerID = "bob"
+			case "tenant":
+				j.TenantID = 2
+			}
+			if err := s.db.Create(&j).Error; err != nil {
+				t.Fatal(err)
+			}
+			ctx = context.WithValue(ctx, leaseContextKey{}, jobLease{j.ID, "lease"})
+			r.Name = "generated"
+			err := s.save(ctx, r, 1)
+			if state == "valid" {
+				if err != nil {
+					t.Fatal(err)
+				}
+			} else if !errors.Is(err, ErrConflict) {
+				t.Fatalf("invalid lease saved: %v", err)
+			}
+			var stored Resource
+			if err := s.db.First(&stored, "id = ?", r.ID).Error; err != nil {
+				t.Fatal(err)
+			}
+			var revisions int64
+			s.db.Model(&Revision{}).Count(&revisions)
+			if state == "valid" && (stored.Version != 2 || revisions != 2) {
+				t.Fatal("valid result must retain the original and create the next revision")
+			}
+			if state != "valid" && (stored.Version != 1 || stored.Name != "original" || revisions != 0) {
+				t.Fatal("failed lease changed persisted state")
+			}
+		})
+	}
+}
+
+func TestManualSaveRejectsActiveJobsAndUnscopedRows(t *testing.T) {
+	s := sqliteWorkbench(t)
+	s.agents = generationAgentStub{}
+	ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
+	r := seedResource(t, s, 1, "alice", "training-maps")
+	r.Data = asJSON(MapData{Topic: "topic", AgentID: "author", DaysCount: 1})
+	if err := s.db.Save(r).Error; err != nil {
+		t.Fatal(err)
+	}
+	if _, err := s.submit(ctx, r.Kind, r.ID, "generate", ActionRequest{Version: 1}); err != nil {
+		t.Fatal(err)
+	}
+	r.Name = "changed"
+	if err := s.save(ctx, r, 1); !errors.Is(err, ErrConflict) {
+		t.Fatalf("queued job must block edit: %v", err)
+	}
+	for _, c := range []context.Context{workbenchContext(1, "bob", types.TenantRoleContributor), workbenchContext(2, "alice", types.TenantRoleAdmin)} {
+		if err := s.save(c, r, 1); err == nil {
+			t.Fatal("cross-owner or cross-tenant write succeeded")
+		}
+		if _, err := s.submit(c, r.Kind, r.ID, "generate", ActionRequest{Version: 1}); err == nil {
+			t.Fatal("cross-owner or cross-tenant submit succeeded")
+		}
+	}
+}
+
+func TestEmptyClaimDoesNotLogRecordNotFound(t *testing.T) {
+	s := sqliteWorkbench(t)
+	var output bytes.Buffer
+	s.db = s.db.Session(&gorm.Session{Logger: logger.New(log.New(&output, "", 0), logger.Config{LogLevel: logger.Warn})})
+	j, err := s.claim(context.Background())
+	if err != nil || j != nil || strings.Contains(output.String(), "record not found") {
+		t.Fatalf("empty queue must be quiet: job=%v err=%v log=%s", j, err, output.String())
+	}
+}
+
+func TestClaimRequiresAffectedRow(t *testing.T) {
+	s := sqliteWorkbench(t)
+	j := Job{ID: uuid.NewString(), TenantID: 1, OwnerID: "alice", ResourceID: uuid.NewString(), Kind: "training-maps", Status: "queued"}
+	if err := s.db.Create(&j).Error; err != nil {
+		t.Fatal(err)
+	}
+	if err := s.db.Exec("CREATE TRIGGER refuse_claim BEFORE UPDATE ON workbench_jobs BEGIN SELECT RAISE(IGNORE); END;").Error; err != nil {
+		t.Fatal(err)
+	}
+	claimed, err := s.claim(context.Background())
+	if claimed != nil || !errors.Is(err, ErrConflict) {
+		t.Fatalf("unclaimed job must not run: job=%v err=%v", claimed, err)
+	}
+}
+
+func TestExpiredLeaseCannotPublishProgress(t *testing.T) {
+	s := sqliteWorkbench(t)
+	until := time.Now().Add(-time.Second)
+	j := Job{ID: uuid.NewString(), Status: "running", LeaseToken: "lease", LeaseUntil: &until}
+	if err := s.db.Create(&j).Error; err != nil {
+		t.Fatal(err)
+	}
+	ctx := context.WithValue(context.Background(), leaseContextKey{}, jobLease{j.ID, j.LeaseToken})
+	if err := s.progress(ctx, j.ID, 90, "done"); !errors.Is(err, ErrConflict) {
+		t.Fatalf("expired progress update allowed: %v", err)
+	}
+}
+
+func TestWorkerRejectsSuspendedMembership(t *testing.T) {
+	s := sqliteWorkbench(t)
+	if err := s.db.AutoMigrate(&types.TenantMember{}); err != nil {
+		t.Fatal(err)
+	}
+	member := types.TenantMember{TenantID: 1, UserID: "alice", Role: types.TenantRoleAdmin, Status: types.TenantMemberStatusSuspended}
+	if err := s.db.Create(&member).Error; err != nil {
+		t.Fatal(err)
+	}
+	r := seedResource(t, s, 1, "alice", "training-maps")
+	until := time.Now().Add(time.Minute)
+	j := Job{ID: uuid.NewString(), TenantID: 1, OwnerID: "alice", ResourceID: r.ID, Kind: r.Kind, Status: "running", LeaseToken: "valid-lease", LeaseUntil: &until, Request: asJSON(ActionRequest{Version: 1})}
+	if err := s.db.Create(&j).Error; err != nil {
+		t.Fatal(err)
+	}
+	s.run(context.Background(), &j)
+	var stored Job
+	if err := s.db.First(&stored, "id = ?", j.ID).Error; err != nil {
+		t.Fatal(err)
+	}
+	if stored.Status != "failed" || stored.Error != ErrForbidden.Error() {
+		t.Fatalf("suspended user executed job: %+v", stored)
+	}
+}
+
+func TestResourceJobListIsScopedAndDoesNotExposeLease(t *testing.T) {
+	s := sqliteWorkbench(t)
+	r := seedResource(t, s, 1, "alice", "training-maps")
+	jobs := []Job{
+		{ID: uuid.NewString(), TenantID: 1, OwnerID: "alice", ResourceID: r.ID, Kind: r.Kind, Status: "running", LeaseToken: "private-token", Request: asJSON(ActionRequest{Advice: "private-advice"})},
+		{ID: uuid.NewString(), TenantID: 1, OwnerID: "bob", ResourceID: r.ID, Kind: r.Kind, Status: "queued"},
+		{ID: uuid.NewString(), TenantID: 2, OwnerID: "alice", ResourceID: r.ID, Kind: r.Kind, Status: "queued"},
+	}
+	if err := s.db.Create(&jobs).Error; err != nil {
+		t.Fatal(err)
+	}
+	for _, role := range []types.TenantRole{types.TenantRoleContributor, types.TenantRoleAdmin} {
+		w := httptest.NewRecorder()
+		c, _ := gin.CreateTestContext(w)
+		c.Params = gin.Params{{Key: "kind", Value: r.Kind}, {Key: "id", Value: r.ID}}
+		c.Request = httptest.NewRequest(http.MethodGet, "/", nil).WithContext(workbenchContext(1, "alice", role))
+		NewHandler(s).Jobs(c)
+		var result struct {
+			Total int64 `json:"total"`
+			Data  []Job `json:"data"`
+		}
+		if err := json.Unmarshal(w.Body.Bytes(), &result); err != nil {
+			t.Fatal(err)
+		}
+		want := int64(1)
+		if role == types.TenantRoleAdmin {
+			want = 2
+		}
+		if w.Code != 200 || result.Total != want || int64(len(result.Data)) != want || strings.Contains(w.Body.String(), "private-token") || strings.Contains(w.Body.String(), "private-advice") {
+			t.Fatalf("job list leaked scope or internal fields: %d %s", w.Code, w.Body.String())
+		}
+	}
+}
+
+func deleteWorkbenchResource(s *Service, ctx context.Context, r *Resource, detach bool) *httptest.ResponseRecorder {
+	gin.SetMode(gin.TestMode)
+	w := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(w)
+	c.Params = gin.Params{{Key: "kind", Value: r.Kind}, {Key: "id", Value: r.ID}}
+	c.Request = httptest.NewRequest(http.MethodDelete, "/?detach="+map[bool]string{true: "true", false: "false"}[detach], nil).WithContext(ctx)
+	NewHandler(s).Delete(c)
+	return w
+}
+
+func TestDeleteBlocksQueuedJobAndDetachesWithRevision(t *testing.T) {
+	s := sqliteWorkbench(t)
+	ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
+	r := seedResource(t, s, 1, "alice", "training-maps")
+	r.Status = "confirmed"
+	r.Data = asJSON(MapData{Days: []Day{{Day: 1, Title: "day"}}})
+	if err := s.db.Save(r).Error; err != nil {
+		t.Fatal(err)
+	}
+	flow := seedResource(t, s, 1, "alice", "flows")
+	flow.Data = asJSON(FlowData{Nodes: []FlowNode{{ID: "n", AgentID: "author"}}})
+	if err := s.db.Save(flow).Error; err != nil {
+		t.Fatal(err)
+	}
+	s.agents = generationAgentStub{}
+	task := seedResource(t, s, 1, "alice", "content-tasks")
+	task.Data = asJSON(TaskData{MapID: r.ID, FlowID: flow.ID, Day: 1, AgentID: "author", Review: "old review", Composed: "old draft"})
+	if err := s.db.Model(task).Update("data", task.Data).Error; err != nil {
+		t.Fatal(err)
+	}
+	if _, err := s.submit(ctx, task.Kind, task.ID, "decompose", ActionRequest{Version: 1}); err != nil {
+		t.Fatal(err)
+	}
+	if w := deleteWorkbenchResource(s, ctx, r, true); w.Code != 409 {
+		t.Fatalf("active dependent must block delete: %d %s", w.Code, w.Body.String())
+	}
+	if err := s.db.Model(&Job{}).Where("resource_id = ?", task.ID).Update("status", "failed").Error; err != nil {
+		t.Fatal(err)
+	}
+	if w := deleteWorkbenchResource(s, ctx, r, false); w.Code != 409 {
+		t.Fatal("delete requires explicit detach")
+	}
+	if w := deleteWorkbenchResource(s, ctx, r, true); w.Code != 200 {
+		t.Fatalf("detach delete: %d %s", w.Code, w.Body.String())
+	}
+	stored, err := s.get(ctx, task.Kind, task.ID)
+	if err != nil {
+		t.Fatal(err)
+	}
+	d, _ := decode[TaskData](stored.Data)
+	var revisions int64
+	s.db.Model(&Revision{}).Where("resource_id = ?", task.ID).Count(&revisions)
+	if stored.Version != 2 || stored.Status != "needs_configuration" || d.MapID != "" || d.Review != "old review" || d.Composed != "old draft" || revisions != 1 {
+		t.Fatal("detachment must preserve historical output and record a revision")
+	}
+}
+
+func TestManualUpdateCannotInjectStateAndReassignmentInvalidatesFollowing(t *testing.T) {
+	original := TaskData{Segments: []Segment{{ID: "s1", AgentID: "a", Status: "accepted", Content: "first"}, {ID: "s2", AgentID: "a", Status: "accepted", Content: "second"}, {ID: "s3", AgentID: "a", Status: "generated", Content: "third"}}, Review: "review", Composed: "final"}
+	r := &Resource{Kind: "content-tasks", Data: asJSON(original), Status: "completed"}
+	injected, _ := decode[TaskData](asJSON(original))
+	injected.Segments[1].Content = "injected"
+	if err := applyManualUpdate(r, resourceRequest{Name: "task", Data: asJSON(injected)}); !errors.Is(err, ErrInvalid) {
+		t.Fatalf("injected content accepted: %v", err)
+	}
+	changed, _ := decode[TaskData](asJSON(original))
+	changed.Segments[1].AgentID = "b"
+	if err := applyManualUpdate(r, resourceRequest{Name: "task", Data: asJSON(changed)}); err != nil {
+		t.Fatal(err)
+	}
+	actual, _ := decode[TaskData](r.Data)
+	if actual.Segments[0].Content != "first" || actual.Segments[0].Status != "accepted" || actual.Segments[1].AgentID != "b" || actual.Segments[1].Content != "second" || actual.Segments[1].Status != "needs_review" || actual.Segments[2].Status != "needs_review" || actual.Review != "" || actual.Composed != "" {
+		t.Fatalf("reassignment must invalidate its dependent output: %+v", actual)
+	}
+	md := MapData{Topic: "before", DaysCount: 1, Days: []Day{{Day: 1, Confirmed: true}}}
+	r = &Resource{Kind: "training-maps", Data: asJSON(md), Status: "confirmed"}
+	md.Days[0].Title = "injected"
+	if err := applyManualUpdate(r, resourceRequest{Data: asJSON(md)}); !errors.Is(err, ErrInvalid) {
+		t.Fatal("direct map content mutation accepted")
+	}
+	md, _ = decode[MapData](r.Data)
+	md.Topic = "after"
+	if err := applyManualUpdate(r, resourceRequest{Name: "map", Data: asJSON(md)}); err != nil {
+		t.Fatal(err)
+	}
+	md, _ = decode[MapData](r.Data)
+	if len(md.Days) != 1 || md.Days[0].Confirmed || r.Status != "needs_review" {
+		t.Fatal("configuration change must retain output and require its review")
+	}
+}
+
+// Uses only an isolated schema. Set WORKBENCH_TEST_POSTGRES_DSN to opt in.
+func TestPostgresResourceRaces(t *testing.T) {
+	dsn := os.Getenv("WORKBENCH_TEST_POSTGRES_DSN")
+	if dsn == "" {
+		t.Skip("WORKBENCH_TEST_POSTGRES_DSN is not set")
+	}
+	db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
+	if err != nil {
+		t.Fatal(err)
+	}
+	sqlDB, _ := db.DB()
+	defer sqlDB.Close()
+	schema := "wb_test_" + strings.ReplaceAll(uuid.NewString(), "-", "")
+	if err = db.Exec("CREATE SCHEMA " + schema).Error; err != nil {
+		t.Fatal(err)
+	}
+	defer db.Exec("DROP SCHEMA " + schema + " CASCADE")
+	// Every pooled connection uses the isolated schema, including parallel workers.
+	testDSN := dsn + " search_path=" + schema
+	if strings.Contains(dsn, "://") {
+		separator := "?"
+		if strings.Contains(dsn, "?") {
+			separator = "&"
+		}
+		testDSN = dsn + separator + "search_path=" + schema
+	}
+	testDB, err := gorm.Open(postgres.Open(testDSN), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
+	if err != nil {
+		t.Fatal(err)
+	}
+	testSQL, _ := testDB.DB()
+	defer testSQL.Close()
+	if err = testDB.AutoMigrate(&Resource{}, &Job{}, &Revision{}, &Event{}, &Consumption{}); err != nil {
+		t.Fatal(err)
+	}
+	s := NewService(testDB)
+	s.agents = generationAgentStub{}
+	ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
+	seedGeneratableMap := func() *Resource {
+		r := seedResource(t, s, 1, "alice", "training-maps")
+		r.Data = asJSON(MapData{Topic: "topic", AgentID: "author", DaysCount: 1})
+		if err := s.db.Save(r).Error; err != nil {
+			t.Fatal(err)
+		}
+		return r
+	}
+	t.Run("lease expires while waiting for lock", func(t *testing.T) {
+		r := seedResource(t, s, 1, "alice", "training-maps")
+		until := time.Now().Add(150 * time.Millisecond)
+		job := Job{ID: uuid.NewString(), TenantID: 1, OwnerID: "alice", ResourceID: r.ID, Kind: r.Kind, Status: "running", LeaseToken: "locked-lease", LeaseUntil: &until}
+		if err := s.db.Create(&job).Error; err != nil {
+			t.Fatal(err)
+		}
+		locked := s.db.Begin()
+		if locked.Error != nil {
+			t.Fatal(locked.Error)
+		}
+		defer locked.Rollback()
+		if err := locked.Exec("SELECT id FROM workbench_jobs WHERE id = ? FOR UPDATE", job.ID).Error; err != nil {
+			t.Fatal(err)
+		}
+		leaseCtx := context.WithValue(ctx, leaseContextKey{}, jobLease{job.ID, job.LeaseToken})
+		finished := make(chan error, 1)
+		go func() { candidate := *r; candidate.Name = "too late"; finished <- s.save(leaseCtx, &candidate, 1) }()
+		time.Sleep(200 * time.Millisecond)
+		if err := locked.Commit().Error; err != nil {
+			t.Fatal(err)
+		}
+		if err := <-finished; !errors.Is(err, ErrConflict) {
+			t.Fatalf("lease expired under lock was allowed to save: %v", err)
+		}
+	})
+	for i := 0; i < 12; i++ {
+		r := seedGeneratableMap()
+		start := make(chan struct{})
+		var wg sync.WaitGroup
+		var editErr, submitErr error
+		wg.Add(2)
+		go func() {
+			defer wg.Done()
+			<-start
+			candidate := *r
+			candidate.Name = "edited"
+			editErr = s.save(ctx, &candidate, 1)
+		}()
+		go func() {
+			defer wg.Done()
+			<-start
+			_, submitErr = s.submit(ctx, r.Kind, r.ID, "generate", ActionRequest{Version: 1})
+		}()
+		close(start)
+		wg.Wait()
+		if (editErr == nil) == (submitErr == nil) {
+			t.Fatalf("exactly one concurrent write must succeed: edit=%v submit=%v", editErr, submitErr)
+		}
+	}
+	t.Run("binding versus parent delete", func(t *testing.T) {
+		for i := 0; i < 8; i++ {
+			m := seedResource(t, s, 1, "alice", "training-maps")
+			if err := s.db.Model(m).Update("status", "confirmed").Error; err != nil {
+				t.Fatal(err)
+			}
+			flow := seedResource(t, s, 1, "alice", "flows")
+			task := &Resource{ID: uuid.NewString(), TenantID: 1, OwnerID: "alice", Kind: "content-tasks", Name: "bound task", Version: 1, Data: asJSON(TaskData{MapID: m.ID, FlowID: flow.ID})}
+			start := make(chan struct{})
+			var wg sync.WaitGroup
+			var createErr error
+			var deleted *httptest.ResponseRecorder
+			wg.Add(2)
+			go func() { defer wg.Done(); <-start; createErr = s.create(ctx, task) }()
+			go func() { defer wg.Done(); <-start; deleted = deleteWorkbenchResource(s, ctx, m, true) }()
+			close(start)
+			wg.Wait()
+			if deleted.Code != 200 {
+				t.Fatalf("parent deletion failed: %d %s", deleted.Code, deleted.Body.String())
+			}
+			if createErr == nil {
+				stored, err := s.get(ctx, task.Kind, task.ID)
+				if err != nil {
+					t.Fatal(err)
+				}
+				data, _ := decode[TaskData](stored.Data)
+				if data.MapID != "" {
+					t.Fatal("concurrent create left dangling reference")
+				}
+			}
+		}
+	})
+	for i := 0; i < 12; i++ {
+		r := seedGeneratableMap()
+		start := make(chan struct{})
+		var wg sync.WaitGroup
+		var deleted *httptest.ResponseRecorder
+		var submitErr error
+		wg.Add(2)
+		go func() { defer wg.Done(); <-start; deleted = deleteWorkbenchResource(s, ctx, r, true) }()
+		go func() {
+			defer wg.Done()
+			<-start
+			_, submitErr = s.submit(ctx, r.Kind, r.ID, "generate", ActionRequest{Version: 1})
+		}()
+		close(start)
+		wg.Wait()
+		if (deleted.Code == 200) == (submitErr == nil) {
+			t.Fatalf("delete and submit must not both succeed: delete=%d submit=%v", deleted.Code, submitErr)
+		}
+	}
+	r := seedGeneratableMap()
+	if _, err = s.submit(ctx, r.Kind, r.ID, "generate", ActionRequest{Version: 1}); err != nil {
+		t.Fatal(err)
+	}
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+	seen := map[string]bool{}
+	for i := 0; i < 8; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			job, err := s.claim(ctx)
+			if err != nil {
+				t.Errorf("claim: %v", err)
+				return
+			}
+			if job != nil {
+				mu.Lock()
+				if seen[job.ID] {
+					t.Error("same job claimed twice")
+				}
+				seen[job.ID] = true
+				mu.Unlock()
+			}
+		}()
+	}
+	wg.Wait()
+}