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

internal/workbench/agent_fusion_test.go

@@ -0,0 +1,216 @@
+package workbench
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+	"github.com/stretchr/testify/require"
+)
+
+type fusionAgents struct {
+	interfaces.CustomAgentService
+	rows    map[string]*types.CustomAgent
+	blocked map[string]bool
+	lookups []string
+}
+
+func (f *fusionAgents) GetAgentByID(ctx context.Context, id string) (*types.CustomAgent, error) {
+	f.lookups = append(f.lookups, id)
+	if f.blocked[id] || f.rows[id] == nil {
+		return nil, ErrForbidden
+	}
+	a := *f.rows[id]
+	return &a, nil
+}
+
+type fusionInstructions struct{ blocked bool }
+
+func (f *fusionInstructions) ResolveAgentInstructions(_ context.Context, a *types.CustomAgent) (string, error) {
+	if len(a.Config.InstructionSkillIDs) > 0 {
+		if f.blocked {
+			return "", ErrForbidden
+		}
+		return "技能表达-" + a.Config.InstructionSkillIDs[0], nil
+	}
+	return "", nil
+}
+
+func fusionHarness(t *testing.T) (*Service, context.Context, *generationChatStub, *fusionAgents) {
+	s, ctx, _, model := generationHarness(t)
+	a := &fusionAgents{rows: map[string]*types.CustomAgent{}, blocked: map[string]bool{}}
+	for _, id := range []string{"primary", "aux", "nested"} {
+		a.rows[id] = &types.CustomAgent{ID: id, Name: id, TenantID: 1, Visibility: "private", Enabled: true,
+			Config: types.CustomAgentConfig{Role: types.AgentRoleProduction, ModelID: "model", SystemPrompt: id + "-identity", MaxCompletionTokens: 1000}}
+	}
+	a.rows["primary"].Config.FusionAgentIDs = []string{"aux"}
+	a.rows["primary"].Config.KnowledgeBases = []string{"primary-kb"}
+	a.rows["aux"].Config.FusionAgentIDs = []string{"primary", "nested"}
+	a.rows["aux"].Config.KnowledgeBases = []string{"private-aux-kb"}
+	a.rows["aux"].Config.InstructionSkillIDs = []string{"original"}
+	s.agents = a
+	s.instructions = &fusionInstructions{}
+	return s, context.WithValue(ctx, agentSnapshotsKey{}, map[string]AgentSnapshot{}), model, a
+}
+
+func TestFusionFrozenDirectExpressionKeepsPrimaryScopeAndDoesNotRecurse(t *testing.T) {
+	s, ctx, model, agents := fusionHarness(t)
+	require.NoError(t, func() error { _, err := s.complete(ctx, "primary", "任务", "延续已确认正文"); return err }())
+	first := model.messages[0].Content
+	require.Contains(t, first, "primary-identity")
+	require.Contains(t, first, "aux-identity")
+	require.Contains(t, first, "技能表达-original")
+	require.Contains(t, first, "主 Agent 的身份、情绪基调、叙事视角及任务要求优先")
+	require.NotContains(t, first, "private-aux-kb")
+	require.NotContains(t, first, "nested-identity")
+	require.NotContains(t, agents.lookups, "nested")
+	saved, ok := snapshotFromContext(ctx, "primary")
+	require.True(t, ok)
+	require.Equal(t, []string{"primary-kb"}, saved.KnowledgeBases)
+	require.Len(t, saved.Fusions, 1)
+	agents.rows["aux"].Config.SystemPrompt = "new-persona"
+	agents.rows["aux"].Config.InstructionSkillIDs = []string{"new-skill"}
+	agents.rows["primary"].Config.FusionAgentIDs = []string{"nested"}
+	_, err := s.complete(ctx, "primary", "任务", "后续段")
+	require.NoError(t, err)
+	require.Contains(t, model.messages[0].Content, "aux-identity")
+	require.Contains(t, model.messages[0].Content, "技能表达-original")
+	require.NotContains(t, model.messages[0].Content, "new-persona")
+}
+
+func TestFusionRetrievalReceivesOnlyThePrimaryTaskKnowledgeScope(t *testing.T) {
+	s, ctx, _, _ := fusionHarness(t)
+	search := &selectionSearchStub{}
+	s.sessions = search
+	kbs, files := []string{"primary-kb"}, []string{"selected-document"}
+	tags := []types.TagScope{{KnowledgeBaseID: "tag-kb", TagIDs: []string{"selected-tag"}}}
+	_, _, err := s.completeWithEvidence(ctx, "primary", "任务", "正文", kbs, files, "主题", tags)
+	require.NoError(t, err)
+	require.Equal(t, kbs, search.kbs)
+	require.Equal(t, files, search.files)
+	require.Equal(t, tags, search.tags)
+	require.Equal(t, "writer", search.owner)
+}
+
+func TestFusionRechecksPermissionRoleEnabledAndPersonaBeforeModelCall(t *testing.T) {
+	for _, change := range []string{"permission", "tenant", "enabled", "role", "skill", "persona"} {
+		t.Run(change, func(t *testing.T) {
+			s, ctx, model, agents := fusionHarness(t)
+			agents.rows["aux"].Config.SystemPrompt += "<!-- weknora:digital-human:original:start -->persona<!-- weknora:digital-human:original:end -->"
+			c, err := s.prepareCompletion(ctx, "primary", "system", "prompt")
+			require.NoError(t, err)
+			switch change {
+			case "permission":
+				agents.blocked["aux"] = true
+			case "tenant":
+				agents.rows["aux"].TenantID = 2
+			case "enabled":
+				agents.rows["aux"].Enabled = false
+			case "role":
+				agents.rows["aux"].Config.Role = types.AgentRoleQC
+			case "skill":
+				s.instructions.(*fusionInstructions).blocked = true
+			case "persona":
+				agents.rows["aux"].Config.SystemPrompt = "new-human"
+			}
+			_, err = s.runCompletion(ctx, c)
+			require.Error(t, err)
+			require.Empty(t, model.messages)
+		})
+	}
+}
+
+func TestOldFrozenSnapshotDoesNotImportNewLiveFusions(t *testing.T) {
+	s, ctx, model, agents := fusionHarness(t)
+	agents.rows["primary"].Config.FusionAgentIDs = nil
+	_, err := s.complete(ctx, "primary", "system", "prompt")
+	require.NoError(t, err)
+	agents.rows["primary"].Config.FusionAgentIDs = []string{"aux"}
+	_, err = s.complete(ctx, "primary", "system", "prompt")
+	require.NoError(t, err)
+	require.NotContains(t, model.messages[0].Content, "aux-identity")
+}
+
+type fusionMemory struct {
+	interfaces.MemoryService
+	items   []*types.MemoryItem
+	users   []string
+	enabled bool
+}
+
+func (m *fusionMemory) Recall(ctx context.Context, _ string) interfaces.MemoryRecall {
+	user, _ := types.UserIDFromContext(ctx)
+	m.users = append(m.users, user)
+	return interfaces.MemoryRecall{Items: m.items}
+}
+func (m *fusionMemory) MemoryAvailable(context.Context) bool { return m.enabled }
+
+func TestFusionMemoryIsCurrentUserConfirmedFrozenAndRevocable(t *testing.T) {
+	s, ctx, model, _ := fusionHarness(t)
+	require.NoError(t, s.db.AutoMigrate(&types.MemoryItem{}))
+	items := []*types.MemoryItem{
+		{ID: "mine", TenantID: 1, SubjectID: "web_user:writer", Kind: types.MemoryKindPreference, Content: "confirmed-own-memory", Status: types.MemoryStatusActive},
+		{ID: "pending", TenantID: 1, SubjectID: "web_user:writer", Kind: types.MemoryKindPreference, Content: "pending-memory", Status: types.MemoryStatusPending},
+		{ID: "other", TenantID: 1, SubjectID: "web_user:other", Kind: types.MemoryKindPreference, Content: "private-other-memory", Status: types.MemoryStatusActive},
+		{ID: "duplicate", TenantID: 1, SubjectID: "web_user:writer", Kind: types.MemoryKindPreference, Content: "confirmed-own-memory", Status: types.MemoryStatusActive},
+	}
+	for _, item := range items {
+		require.NoError(t, s.db.Create(item).Error)
+	}
+	memory := &fusionMemory{items: items, enabled: true}
+	s.memories = memory
+	_, err := s.complete(ctx, "primary", "system", "prompt")
+	require.NoError(t, err)
+	require.Equal(t, []string{"writer"}, memory.users)
+	require.Equal(t, 1, strings.Count(model.messages[0].Content, "confirmed-own-memory"))
+	require.NotContains(t, model.messages[0].Content, "pending-memory")
+	require.NotContains(t, model.messages[0].Content, "private-other-memory")
+	c, err := s.prepareCompletion(ctx, "primary", "system", "prompt")
+	require.NoError(t, err)
+	require.NoError(t, s.db.Model(&types.MemoryItem{}).Where("id = ?", "mine").Update("content", "edited-after-freeze").Error)
+	_, err = s.runCompletion(ctx, c)
+	require.NoError(t, err)
+	require.NotContains(t, model.messages[0].Content, "confirmed-own-memory")
+	require.NotContains(t, model.messages[0].Content, "edited-after-freeze")
+	require.Len(t, memory.users, 1)
+	other := context.WithValue(ctx, types.UserIDContextKey, "other")
+	otherRequest, err := s.prepareCompletion(other, "primary", "system", "prompt")
+	require.NoError(t, err)
+	require.NotContains(t, otherRequest.messages[0].Content, "confirmed-own-memory")
+	require.Empty(t, s.currentFusionMemories(other, items[:1]))
+}
+
+func TestFusionMemoryHonorsLiveAgentAndUserOptOut(t *testing.T) {
+	for _, scope := range []string{"primary", "aux", "user"} {
+		t.Run(scope, func(t *testing.T) {
+			s, ctx, model, agents := fusionHarness(t)
+			require.NoError(t, s.db.AutoMigrate(&types.MemoryItem{}))
+			item := &types.MemoryItem{ID: "mine", TenantID: 1, SubjectID: "web_user:writer", Kind: types.MemoryKindPreference, Content: "confirmed-memory", Status: types.MemoryStatusActive}
+			require.NoError(t, s.db.Create(item).Error)
+			memory := &fusionMemory{items: []*types.MemoryItem{item}, enabled: true}
+			s.memories = memory
+			c, err := s.prepareCompletion(ctx, "primary", "system", "prompt")
+			require.NoError(t, err)
+			require.Contains(t, c.messages[0].Content, item.Content)
+			if scope == "user" {
+				memory.enabled = false
+			} else {
+				disabled := false
+				agents.rows[scope].Config.MemoryEnabled = &disabled
+			}
+			_, err = s.runCompletion(ctx, c)
+			require.NoError(t, err)
+			require.NotContains(t, model.messages[0].Content, item.Content)
+		})
+	}
+}
+
+func TestFusionPromptHasGlobalBudget(t *testing.T) {
+	fusions := make([]FusionSnapshot, types.MaxFusionAgents)
+	for i := range fusions {
+		fusions[i] = FusionSnapshot{Name: strings.Repeat("名", 1000), SystemPrompt: strings.Repeat("风格", 10000), Instructions: strings.Repeat("技巧", 10000)}
+	}
+	require.LessOrEqual(t, len([]rune(fusionPrompt(fusions, nil))), fusionPromptRuneLimit)
+}