xinmi/internal/workbench/agent_fusion_test.go

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)
}