xinmi/internal/workbench/knowledge_selection_test.go

package workbench

import (
	"context"
	"reflect"
	"testing"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
)

func TestMapKnowledgeFoldersSurviveConfirmation(t *testing.T) {
	folders := map[string][]string{"kb-one": {"folder-a", "folder-b"}, "kb-two": {"folder-c"}}
	tags := map[string][]string{"kb-one": {"tag-a"}}
	d := MapData{Topic: "topic", DaysCount: 1, AgentID: "agent", KnowledgeIDs: []string{"file-a"}, KnowledgeFolders: folders, KnowledgeTags: tags, Days: []Day{{Day: 1, Title: "Foundation", Hook: "Next", Nodes: []Node{{ID: "n1", Label: "Concept"}}, Framework: "Confirmed text"}}}
	r := &Resource{Kind: "training-maps", Data: asJSON(d)}
	s := NewService(nil)
	if err := s.executeMap(context.Background(), &Job{Action: "confirm-day"}, r, ActionRequest{Day: 1}); err != nil {
		t.Fatal(err)
	}
	stored, err := decode[MapData](r.Data)
	if err != nil || !stored.Days[0].Confirmed || !reflect.DeepEqual(stored.KnowledgeFolders, folders) || !reflect.DeepEqual(stored.KnowledgeTags, tags) || !reflect.DeepEqual(stored.KnowledgeIDs, d.KnowledgeIDs) {
		t.Fatalf("confirmation dropped the selected categories: %+v, %v", stored, err)
	}
}

func TestMapKnowledgeFoldersSurviveAIGeneration(t *testing.T) {
	s, ctx, job, model := generationHarness(t)
	job.Action = "generate"
	model.output = `{"days":[{"day":1,"title":"Foundation","hook":"Next step","nodes":[{"id":"n1","label":"Concept"}]}]}`
	folders := map[string][]string{"kb-one": {"folder-a"}}
	r := &Resource{Kind: "training-maps", Data: asJSON(MapData{Topic: "Topic", DaysCount: 1, AgentID: "agent", KnowledgeFolders: folders})}
	if err := s.executeMap(ctx, job, r, ActionRequest{}); err != nil {
		t.Fatal(err)
	}
	stored, err := decode[MapData](r.Data)
	if err != nil || len(stored.Days) != 1 || !reflect.DeepEqual(stored.KnowledgeFolders, folders) {
		t.Fatalf("generation dropped the selected categories: %+v, %v", stored, err)
	}
}

func TestTaskKnowledgeFoldersSurviveGeneratedStateUpdates(t *testing.T) {
	folders := map[string][]string{"kb-one": {"folder-a"}}
	d := TaskData{KnowledgeFolders: folders, KnowledgeIDs: []string{"file-a"}, Segments: []Segment{{ID: "s1", AgentID: "agent", Content: "old", Status: "generated"}}}
	r := &Resource{Kind: "content-tasks", Name: "task", Data: asJSON(d)}
	s := NewService(nil)
	if err := s.executeTask(context.Background(), &Job{Action: "skip"}, r, ActionRequest{SegmentID: "s1"}); err != nil {
		t.Fatal(err)
	}
	stored, err := decode[TaskData](r.Data)
	if err != nil || stored.Segments[0].Status != "skipped" || !reflect.DeepEqual(stored.KnowledgeFolders, folders) {
		t.Fatalf("task action dropped categories: %+v, %v", stored, err)
	}
	if err = applyManualUpdate(r, resourceRequest{Name: "renamed", Data: r.Data}); err != nil {
		t.Fatal(err)
	}
	stored, err = decode[TaskData](r.Data)
	if err != nil || !reflect.DeepEqual(stored.KnowledgeFolders, folders) {
		t.Fatalf("manual update dropped categories: %+v, %v", stored, err)
	}
}

type selectionSearchStub struct {
	interfaces.SessionService
	tenant     uint64
	owner      string
	role       types.TenantRole
	kbs, files []string
	tags       []types.TagScope
}

func (s *selectionSearchStub) SearchKnowledge(ctx context.Context, kbs, files []string, tags []types.TagScope, _ string) ([]*types.SearchResult, error) {
	s.tenant, _ = types.TenantIDFromContext(ctx)
	s.owner = types.SessionOwnerIDFromContext(ctx)
	s.role = types.TenantRoleFromContext(ctx)
	s.kbs, s.files = append([]string(nil), kbs...), append([]string(nil), files...)
	s.tags = tags
	return []*types.SearchResult{{ID: "chunk", KnowledgeID: "folder-only-file", Content: "evidence"}}, nil
}

func TestKnowledgeSelectionUsesMapScopeOrAssignedTaskAgent(t *testing.T) {
	for _, action := range []string{"map-plan", "map-day", "task-decompose", "task-segment"} {
		t.Run(action, func(t *testing.T) {
			s, ctx, job, model := generationHarness(t)
			if err := s.db.Exec("CREATE TABLE knowledges (id TEXT PRIMARY KEY, knowledge_base_id TEXT, folder_path TEXT)").Error; err != nil {
				t.Fatal(err)
			}
			if err := s.db.Exec("INSERT INTO knowledges VALUES ('folder-only-file','mixed-kb','folder-a')").Error; err != nil {
				t.Fatal(err)
			}
			search := &selectionSearchStub{}
			s.sessions = search
			folders := map[string][]string{"mixed-kb": {"folder-a"}, "tag-only-kb": {}}
			tags := map[string][]string{"mixed-kb": {"mixed-tag"}, "tag-only-kb": {"other-tag"}}
			files := []string{"folder-only-file", "mixed-tag-file", "other-tag-file"}
			mapData := MapData{Topic: "topic", DaysCount: 1, AgentID: "agent", KnowledgeIDs: files, KnowledgeFolders: folders, KnowledgeTags: tags, Days: []Day{{Day: 1, Title: "Foundation", Hook: "Next", Nodes: []Node{{ID: "n1", Label: "Concept"}}}}}
			resource := &Resource{Kind: "training-maps", Data: asJSON(mapData)}
			request := ActionRequest{Day: 1}
			job.Action = "generate"
			model.output = `{"days":[{"day":1,"title":"Foundation","hook":"Next","nodes":[{"id":"n1","label":"Concept"}]}]}`
			if action == "map-day" {
				job.Action = "generate-day"
				model.output = `{"framework":"# Foundation"}`
			}
			var err error
			if action == "task-decompose" || action == "task-segment" {
				s.agents = generationAgentStub{config: types.CustomAgentConfig{KBSelectionMode: "selected", KnowledgeBases: []string{"agent-kb"}}}
				flow := &Resource{ID: "flow", TenantID: 1, OwnerID: "writer", Kind: "flows", Name: "Flow", Version: 1, Data: asJSON(FlowData{Nodes: []FlowNode{{ID: "n1", AgentID: "agent"}}, KnowledgeIDs: files, KnowledgeFolders: folders, KnowledgeTags: tags})}
				if err = s.db.Create(flow).Error; err != nil {
					t.Fatal(err)
				}
				resource = &Resource{Kind: "content-tasks", Name: "Task", Data: asJSON(TaskData{FlowID: flow.ID, SourceText: "# Foundation", AgentID: "agent", Segments: []Segment{{ID: "s1", Title: "Foundation", AgentID: "agent", Status: "pending"}}})}
				request = ActionRequest{SegmentID: "s1"}
				model.output = "Generated segment"
				if action == "task-decompose" {
					job.Action = "decompose"
					model.output = `{"segments":[{"id":"s1","title":"Foundation","agent_id":"agent"}]}`
				}
				err = s.executeTask(ctx, job, resource, request)
			} else {
				err = s.executeMap(ctx, job, resource, request)
			}
			if err != nil {
				t.Fatal(err)
			}
			if resource.Kind == "content-tasks" {
				if !reflect.DeepEqual(search.kbs, []string{"agent-kb"}) || len(search.files) != 0 || len(search.tags) != 0 {
					t.Fatalf("legacy canvas selection overrode the assigned Agent: kbs=%v files=%v tags=%v", search.kbs, search.files, search.tags)
				}
				if action == "task-decompose" {
					stored, err := decode[TaskData](resource.Data)
					if err != nil || len(stored.KnowledgeIDs) != 0 || len(stored.KnowledgeTags) != 0 || len(stored.KnowledgeFolders) != 0 || stored.FlowSnapshot == nil || len(stored.FlowSnapshot.KnowledgeIDs) != 0 {
						t.Fatalf("legacy canvas scope survived task decomposition: %+v, %v", stored, err)
					}
				}
				return
			}
			if len(search.kbs) != 0 || !reflect.DeepEqual(search.files, files) {
				t.Fatalf("resolved selection changed: kbs=%v files=%v", search.kbs, search.files)
			}
			wantTags := []types.TagScope{{KnowledgeBaseID: "tag-only-kb", TagIDs: []string{"other-tag"}}}
			if !reflect.DeepEqual(search.tags, wantTags) {
				t.Fatalf("mixed selection was re-filtered or independent tag scope lost: %+v", search.tags)
			}
			if compact(resource.Data["knowledge_tags"]) != compact(tags) || compact(resource.Data["knowledge_folders"]) != compact(folders) {
				t.Fatal("generated update dropped selection metadata")
			}
		})
	}
}

func TestEvidencePreservesBackgroundUserScopeAndFileOnlySelection(t *testing.T) {
	search := &selectionSearchStub{}
	s := NewService(nil)
	s.sessions = search
	d := MapData{KnowledgeIDs: []string{"allowed-file"}, KnowledgeFolders: map[string][]string{"metadata-only-kb": {"metadata-only-folder"}}}
	ctx := types.WithBackgroundTask(workbenchContext(4, "alice", types.TenantRoleContributor))
	if _, _, err := s.evidence(ctx, d.KBIDs, d.KnowledgeIDs, "query"); err != nil {
		t.Fatal(err)
	}
	if search.owner != "alice" || search.tenant != 4 || search.role != types.TenantRoleContributor {
		t.Fatalf("retrieval lost worker identity: %+v", search)
	}
	if len(search.kbs) != 0 || !reflect.DeepEqual(search.files, []string{"allowed-file"}) {
		t.Fatalf("folder metadata must not broaden retrieval: kbs=%v files=%v", search.kbs, search.files)
	}
}