212 lines
8.7 KiB
Go
212 lines
8.7 KiB
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
|
||
|
|
}
|