xinmi/internal/workbench/flow_knowledge_test.go

package workbench

import (
	"context"
	"reflect"
	"testing"

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

type flowKnowledgeBases struct {
	interfaces.KnowledgeBaseService
	ids []string
}

func (s flowKnowledgeBases) GetKnowledgeBaseByID(_ context.Context, id string) (*types.KnowledgeBase, error) {
	for _, allowed := range s.ids {
		if id == allowed {
			return &types.KnowledgeBase{ID: id, TenantID: 1}, nil
		}
	}
	return nil, ErrForbidden
}

func (s flowKnowledgeBases) ListKnowledgeBases(context.Context) ([]*types.KnowledgeBase, error) {
	var bases []*types.KnowledgeBase
	for _, id := range s.ids {
		bases = append(bases, &types.KnowledgeBase{ID: id, TenantID: 1})
	}
	return bases, nil
}

type flowKnowledgeFiles struct{ interfaces.KnowledgeService }

func (flowKnowledgeFiles) GetKnowledgeBatchWithSharedAccess(_ context.Context, _ uint64, ids []string) ([]*types.Knowledge, error) {
	if len(ids) > 0 {
		return nil, ErrForbidden
	}
	return nil, nil
}

func TestFlowPersistenceDropsLegacyKnowledgeAndPreservesMetadata(t *testing.T) {
	for _, operation := range []string{"create", "save_supplied_scope", "save_omitted_scope", "rename"} {
		t.Run(operation, func(t *testing.T) {
			s := sqliteWorkbench(t)
			s.agents = generationAgentStub{}
			ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
			data := asJSON(FlowData{Nodes: []FlowNode{{ID: "n", AgentID: "writer"}}, KBIDs: []string{"deleted"}, KnowledgeIDs: []string{"deleted-file"}, KnowledgeFolders: map[string][]string{"deleted": {"folder"}}, KnowledgeTags: map[string][]string{"deleted": {"tag"}}})
			data["opaque_metadata"] = JSON{"nested": []any{"preserved", 42}, "enabled": true}
			data["layout_metadata"] = JSON{"zoom": 0.75}
			r := &Resource{ID: "flow", TenantID: 1, OwnerID: "alice", Kind: "flows", Name: "Original", Status: "draft", Version: 1, Data: data}
			original := compact(data)
			want := asJSON(data)
			if operation == "create" {
				if err := s.validate(ctx, r); err != nil {
					t.Fatal(err)
				}
				if err := s.create(ctx, r); err != nil {
					t.Fatal(err)
				}
			} else {
				if err := s.db.Create(r).Error; err != nil {
					t.Fatal(err)
				}
				proposed := asJSON(data)
				if operation != "rename" {
					proposed["nodes"] = []FlowNode{{ID: "n", Label: "Updated", AgentID: "writer"}}
				}
				if operation == "save_omitted_scope" {
					for _, key := range []string{"knowledge_base_ids", "knowledge_ids", "knowledge_folders", "knowledge_tags"} {
						delete(proposed, key)
					}
				}
				want = asJSON(proposed)
				if response := updateWorkbenchResource(s, ctx, r, resourceRequest{Name: "Updated", Version: 1, Data: proposed}); response.Code != 200 {
					t.Fatal(response.Code, response.Body.String())
				}
			}
			for _, key := range []string{"knowledge_base_ids", "knowledge_ids", "knowledge_folders", "knowledge_tags"} {
				delete(want, key)
			}
			stored, err := s.get(ctx, "flows", r.ID)
			if err != nil {
				t.Fatal(err)
			}
			if compact(stored.Data) != compact(want) {
				t.Fatalf("persisted canvas retained legacy scope or changed other metadata: got %s, want %s", compact(stored.Data), compact(want))
			}
			if operation != "create" {
				var before, after Revision
				if err := s.db.First(&before, "resource_id = ? AND version = 1", r.ID).Error; err != nil {
					t.Fatal(err)
				}
				if err := s.db.First(&after, "resource_id = ? AND version = 2", r.ID).Error; err != nil {
					t.Fatal(err)
				}
				if stored.Version != 2 || compact(before.Data) != original || compact(after.Data) != compact(want) {
					t.Fatal("canvas save lost its original revision or recorded unsanitized data")
				}
			}
		})
	}
}

func TestFlowKnowledgeModesIgnoreStaleCanvasAndTaskScopes(t *testing.T) {
	for _, mode := range []string{"selected", "all", "none"} {
		t.Run(mode, func(t *testing.T) {
			s := NewService(nil)
			s.agents = generationAgentStub{config: types.CustomAgentConfig{KBSelectionMode: mode, KnowledgeBases: []string{"selected"}}}
			s.knowledgeBases = flowKnowledgeBases{ids: []string{"selected", "other"}}
			s.knowledge = flowKnowledgeFiles{}
			flow := FlowData{Nodes: []FlowNode{{ID: "n", AgentID: "agent"}}, KBIDs: []string{"deleted"}, KnowledgeIDs: []string{"deleted-file"}, KnowledgeFolders: map[string][]string{"deleted": {"empty-folder"}}, KnowledgeTags: map[string][]string{"deleted": {"deleted-tag"}}}
			d := TaskData{KBIDs: flow.KBIDs, KnowledgeIDs: flow.KnowledgeIDs, KnowledgeFolders: flow.KnowledgeFolders, KnowledgeTags: flow.KnowledgeTags}
			if err := s.inheritFlow(workbenchContext(1, "writer", types.TenantRoleContributor), &d, flow); err != nil {
				t.Fatal(err)
			}
			var want []string
			if mode == "selected" {
				want = []string{"selected"}
			} else if mode == "all" {
				want = []string{"selected", "other"}
			}
			if !reflect.DeepEqual(d.KBIDs, want) || len(d.KnowledgeIDs) != 0 || len(d.KnowledgeFolders) != 0 || len(d.KnowledgeTags) != 0 {
				t.Fatalf("unexpected agent scope: %+v", d)
			}
		})
	}
}

func TestTaskAgentCategoriesAndGenerationIgnoreLegacyFlowKnowledge(t *testing.T) {
	s, ctx, job, _ := generationHarness(t)
	s.agents = &fusionAgents{rows: map[string]*types.CustomAgent{
		"writer": {ID: "writer", TenantID: 1, Config: types.CustomAgentConfig{Role: types.AgentRoleProduction, ModelID: "model", KBSelectionMode: "selected", KnowledgeBases: []string{"writer-live"}}},
		"editor": {ID: "editor", TenantID: 1, Config: types.CustomAgentConfig{Role: types.AgentRoleProduction, ModelID: "model", KBSelectionMode: "selected", KnowledgeBases: []string{"editor-kb"}}},
	}}
	s.knowledgeBases = flowKnowledgeBases{ids: []string{"writer-live", "writer-frozen", "editor-kb"}}
	s.knowledge = flowKnowledgeFiles{}
	if err := s.db.AutoMigrate(&types.KnowledgeTag{}); err != nil {
		t.Fatal(err)
	}
	for _, baseID := range []string{"writer-live", "writer-frozen", "editor-kb"} {
		if err := s.db.Create(&types.KnowledgeTag{ID: baseID + "-category", TenantID: 1, KnowledgeBaseID: baseID, Name: baseID}).Error; err != nil {
			t.Fatal(err)
		}
	}
	flow := FlowData{Nodes: []FlowNode{{ID: "n1", AgentID: "writer"}, {ID: "n2", AgentID: "editor"}}, KBIDs: []string{"deleted"}, KnowledgeIDs: []string{"deleted-file"}, KnowledgeFolders: map[string][]string{"deleted": {"empty-folder"}}, KnowledgeTags: map[string][]string{"deleted": {"deleted-category"}}}
	r := &Resource{ID: "flow", TenantID: 1, OwnerID: "writer", Kind: "flows", Name: "Flow", Version: 1, Data: asJSON(flow)}
	if err := s.validate(ctx, r); err != nil {
		t.Fatalf("legacy scope prevented saving a valid flow: %v", err)
	}
	if err := s.db.Create(r).Error; err != nil {
		t.Fatal(err)
	}
	d := TaskData{FlowID: r.ID, FlowSnapshot: &flow, AgentSnapshots: map[string]AgentSnapshot{"writer": {ModelID: "model", KBSelectionMode: "selected", KnowledgeBases: []string{"writer-frozen"}}}}
	resolved, err := s.taskFlow(ctx, d)
	if err != nil || len(resolved.KBIDs) != 0 || len(resolved.KnowledgeIDs) != 0 || len(resolved.KnowledgeFolders) != 0 || len(resolved.KnowledgeTags) != 0 {
		t.Fatalf("legacy scope obstructed the flow snapshot: %+v, %v", resolved, err)
	}
	if err := s.inheritFlow(ctx, &d, resolved); err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(d.KBIDs, []string{"writer-frozen", "editor-kb"}) {
		t.Fatalf("planning scope did not retain each selected Agent's snapshot: %v", d.KBIDs)
	}
	options, err := s.agentCategoryOptions(ctx, d, resolved)
	if err != nil {
		t.Fatal(err)
	}
	if len(options["writer"]) != 1 || options["writer"][0].ID != "writer-frozen-category" || len(options["editor"]) != 1 || options["editor"][0].ID != "editor-kb-category" {
		t.Fatalf("categories leaked between Agents or ignored the snapshot: %+v", options)
	}
	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','editor-kb','')").Error; err != nil {
		t.Fatal(err)
	}
	search := &selectionSearchStub{}
	s.sessions = search
	d.Segments = []Segment{{ID: "s1", Title: "Draft", AgentID: "writer", CategoryID: "editor-kb-category", Status: "pending"}}
	resource := &Resource{Kind: "content-tasks", Name: "Task", Data: asJSON(d)}
	job.Action = "generate"
	if err := s.executeTask(ctx, job, resource, ActionRequest{SegmentID: "s1", AgentID: "editor"}); err != nil {
		t.Fatal(err)
	}
	wantTags := []types.TagScope{{KnowledgeBaseID: "editor-kb", TagIDs: []string{"editor-kb-category"}}}
	if len(search.kbs) != 0 || len(search.files) != 0 || !reflect.DeepEqual(search.tags, wantTags) {
		t.Fatalf("reassigned segment lost its Agent/category scope: kbs=%v files=%v tags=%v", search.kbs, search.files, search.tags)
	}
}