xinmi/internal/workbench/memory_legacy_test.go

package workbench

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/Tencent/WeKnora/internal/accesspolicy"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/gin-gonic/gin"
	"gorm.io/gorm"
)

func TestLegacyMemoryReconfirmationPreservesSourceAndNativeLink(t *testing.T) {
	s, ctx, row, calls := legacyMemoryHarness(t)
	w := legacyMemoryRequest(s, ctx, row.ID, memoryReviewRequest{Status: "accepted", Content: "先讲结论\n再给案例", ReconfirmLegacy: true, ConfirmPersonalScope: true})
	if w.Code != http.StatusOK {
		t.Fatalf("reconfirmation failed: %d %s", w.Code, w.Body.String())
	}
	var saved PendingMemory
	if err := s.db.First(&saved, "id = ?", row.ID).Error; err != nil {
		t.Fatal(err)
	}
	if saved.Content != row.Content || saved.SourceRef != row.SourceRef || saved.LegacyHistoryID != row.LegacyHistoryID || saved.Status != "accepted" || saved.NativeMemoryID != "native" || saved.ReconfirmedAt == nil || saved.ReconfirmedContent != "先讲结论 再给案例" || saved.ReviewedBy != row.OwnerID {
		t.Fatalf("source/link not preserved: %+v", saved)
	}
	var native types.MemoryItem
	if err := s.db.First(&native, "id = ?", saved.NativeMemoryID).Error; err != nil {
		t.Fatal(err)
	}
	if native.SourceMessageID != row.ID || native.Kind != types.MemoryKindPreference || native.SubjectID != "web_user:alice" {
		t.Fatalf("native provenance or scope lost: %+v", native)
	}
	// A retry, even after explicit deletion, must not recreate or overwrite memory.
	if err := s.db.Delete(&native).Error; err != nil {
		t.Fatal(err)
	}
	if err := s.db.Where("id = ?", row.AgentID).Delete(&types.CustomAgent{}).Error; err != nil {
		t.Fatal(err)
	}
	w = legacyMemoryRequest(s, ctx, row.ID, memoryReviewRequest{Status: "accepted", Content: "不同的重试内容", ReconfirmLegacy: true, ConfirmPersonalScope: true})
	var count int64
	s.db.Model(&types.MemoryItem{}).Count(&count)
	if w.Code != http.StatusOK || *calls != 1 || count != 0 {
		t.Fatalf("retry resurrected memory: http=%d calls=%d count=%d", w.Code, *calls, count)
	}
	if err := s.db.First(&saved, "id = ?", row.ID).Error; err != nil || saved.ReconfirmedContent != "先讲结论 再给案例" {
		t.Fatalf("retry changed confirmation: %+v %v", saved, err)
	}
}

func TestLegacyMemoryReconfirmationGuardsAndRollback(t *testing.T) {
	for _, mode := range []string{"owner", "tenant", "permission", "no_consent", "empty_content", "long_content", "not_legacy", "pending", "rejected", "no_flag", "reject_request", "deleted_agent", "agent_disabled", "agent_paused", "memory_disabled", "write_failure", "save_failure", "wrong_native_scope", "nil_native"} {
		t.Run(mode, func(t *testing.T) {
			s, ctx, row, calls := legacyMemoryHarness(t)
			body := memoryReviewRequest{Status: "accepted", Content: "先讲结论再给案例", ReconfirmLegacy: true, ConfirmPersonalScope: true}
			switch mode {
			case "owner":
				ctx = workbenchContext(1, "bob", types.TenantRoleAdmin)
			case "tenant":
				ctx = workbenchContext(2, "alice", types.TenantRoleAdmin)
			case "permission":
				ctx = accesspolicy.WithSnapshot(ctx, &accesspolicy.Snapshot{Restricted: true})
			case "no_consent":
				body.ConfirmPersonalScope = false
			case "empty_content":
				body.Content = "  "
			case "long_content":
				body.Content = strings.Repeat("字", 301)
			case "not_legacy":
				row.LegacyHistoryID = ""
			case "pending", "rejected":
				row.Status = mode
			case "no_flag":
				body.ReconfirmLegacy = false
			case "reject_request":
				body.Status = "rejected"
			case "deleted_agent":
				if err := s.db.Where("id = ?", row.AgentID).Delete(&types.CustomAgent{}).Error; err != nil {
					t.Fatal(err)
				}
			case "agent_disabled":
				disabled := false
				agent := types.CustomAgent{ID: row.AgentID, TenantID: 1, Enabled: true, Config: types.CustomAgentConfig{MemoryEnabled: &disabled}}
				if err := s.db.Save(&agent).Error; err != nil {
					t.Fatal(err)
				}
			case "agent_paused":
				if err := s.db.Model(&types.CustomAgent{}).Where("id = ?", row.AgentID).Update("enabled", false).Error; err != nil {
					t.Fatal(err)
				}
			case "memory_disabled":
				s.memories = feedbackMemoryStub{available: false}
			case "write_failure", "wrong_native_scope", "nil_native":
				s.SetTransactionalMemoryWriter(func(ctx context.Context, tx *gorm.DB, item types.MemoryItem) (*types.MemoryItem, error) {
					*calls++
					item.ID, item.TenantID, item.SubjectID, item.Status = "native", 1, "web_user:alice", types.MemoryStatusActive
					if mode == "wrong_native_scope" {
						item.SubjectID = "web_user:bob"
					}
					if err := tx.Create(&item).Error; err != nil {
						return nil, err
					}
					if mode == "write_failure" {
						return nil, errors.New("failed write")
					}
					if mode == "nil_native" {
						return nil, nil
					}
					return &item, nil
				})
			}
			if err := s.db.Save(&row).Error; err != nil {
				t.Fatal(err)
			}
			if mode == "save_failure" {
				if err := s.db.Exec("CREATE TRIGGER reject_reconfirmation BEFORE UPDATE ON workbench_pending_memories BEGIN SELECT RAISE(ABORT, 'failed link'); END").Error; err != nil {
					t.Fatal(err)
				}
			}
			w := legacyMemoryRequest(s, ctx, row.ID, body)
			var saved PendingMemory
			if err := s.db.First(&saved, "id = ?", row.ID).Error; err != nil {
				t.Fatal(err)
			}
			var count int64
			s.db.Model(&types.MemoryItem{}).Count(&count)
			if w.Code < 400 || saved.NativeMemoryID != "" || saved.ReconfirmedAt != nil || saved.Content != row.Content || saved.Status != row.Status || count != 0 {
				t.Fatalf("guard/rollback failed: mode=%s http=%d saved=%+v native=%d", mode, w.Code, saved, count)
			}
		})
	}
}

func TestLegacyMemoryListEligibility(t *testing.T) {
	s, ctx, row, _ := legacyMemoryHarness(t)
	rows := []PendingMemory{row, {ID: "native-history", Status: "accepted"}, {ID: "pending", Status: "pending", LegacyHistoryID: "history"}}
	s.decorateLegacyMemories(ctx, rows)
	if !rows[0].CanReconfirm || rows[1].CanReconfirm || rows[2].CanReconfirm {
		t.Fatalf("incorrect eligibility: %+v", rows)
	}
	rows = []PendingMemory{row}
	s.decorateLegacyMemories(accesspolicy.WithSnapshot(ctx, &accesspolicy.Snapshot{Restricted: true}), rows)
	if rows[0].CanReconfirm || rows[0].ReconfirmationUnavailableReason != "没有记忆管理权限" {
		t.Fatalf("missing permission offered reconfirmation: %+v", rows[0])
	}
	s.memories = feedbackMemoryStub{available: false}
	rows = []PendingMemory{row}
	s.decorateLegacyMemories(ctx, rows)
	if rows[0].CanReconfirm || rows[0].ReconfirmationUnavailableReason == "" {
		t.Fatalf("disabled memory offered reconfirmation: %+v", rows[0])
	}
}

func legacyMemoryHarness(t *testing.T) (*Service, context.Context, PendingMemory, *int) {
	t.Helper()
	s := sqliteWorkbench(t)
	return legacyMemoryHarnessWithService(t, s)
}

func legacyMemoryHarnessWithService(t *testing.T, s *Service) (*Service, context.Context, PendingMemory, *int) {
	t.Helper()
	if err := s.db.AutoMigrate(&types.MemoryItem{}, &types.CustomAgent{}); err != nil {
		t.Fatal(err)
	}
	s.agents = legacyMemoryAgentStub{db: s.db}
	s.memories = feedbackMemoryStub{available: true}
	row := PendingMemory{ID: "legacy-candidate", TenantID: 1, OwnerID: "alice", AgentID: "author", Content: "旧任务中适用的原始建议", SourceRef: "legacy-task/1", Status: "accepted", LegacyHistoryID: "history"}
	if err := s.db.Create(&row).Error; err != nil {
		t.Fatal(err)
	}
	if err := s.db.Create(&types.CustomAgent{ID: row.AgentID, TenantID: 1, CreatedBy: row.OwnerID, Name: "Author", Enabled: true}).Error; err != nil {
		t.Fatal(err)
	}
	calls := 0
	s.SetTransactionalMemoryWriter(func(ctx context.Context, tx *gorm.DB, item types.MemoryItem) (*types.MemoryItem, error) {
		calls++
		item.ID, item.TenantID, item.SubjectID, item.Status = "native", 1, "web_user:alice", types.MemoryStatusActive
		item.Content = types.SanitizeMemoryContent(item.Content)
		return &item, tx.Create(&item).Error
	})
	return s, workbenchContext(1, "alice", types.TenantRoleContributor), row, &calls
}

type legacyMemoryAgentStub struct {
	generationAgentStub
	db *gorm.DB
}

func (a legacyMemoryAgentStub) GetAgentByID(ctx context.Context, id string) (*types.CustomAgent, error) {
	tenant, _ := types.TenantIDFromContext(ctx)
	var agent types.CustomAgent
	err := a.db.WithContext(ctx).Where("id = ? AND tenant_id = ?", id, tenant).First(&agent).Error
	return &agent, err
}

func legacyMemoryRequest(s *Service, ctx context.Context, id string, body memoryReviewRequest) *httptest.ResponseRecorder {
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Params = gin.Params{{Key: "id", Value: id}}
	raw, _ := json.Marshal(body)
	c.Request = httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(raw)).WithContext(ctx)
	c.Request.Header.Set("Content-Type", "application/json")
	NewHandler(s).ReviewMemory(c)
	return w
}