package workbench import ( "context" "encoding/base64" "encoding/json" "errors" "net/http" "net/http/httptest" "strings" "testing" "github.com/Tencent/WeKnora/internal/models/vlm" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" "github.com/gin-gonic/gin" ) func updateWorkbenchResource(s *Service, ctx context.Context, r *Resource, request resourceRequest) *httptest.ResponseRecorder { w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) data, _ := json.Marshal(request) c.Request = httptest.NewRequest(http.MethodPut, "/", strings.NewReader(string(data))).WithContext(ctx) c.Request.Header.Set("Content-Type", "application/json") c.Params = gin.Params{{Key: "kind", Value: r.Kind}, {Key: "id", Value: r.ID}} NewHandler(s).Update(c) return w } func TestImportedResourceRenameConfigurationActivationAndSpoofProtection(t *testing.T) { s := sqliteWorkbench(t) s.agents = generationAgentStub{} ctx := workbenchContext(1, "alice", types.TenantRoleContributor) r := seedResource(t, s, 1, "alice", "flows") r.Status = "needs_configuration" r.Data = JSON{"legacy_import": JSON{"read_only": true, "needs_configuration": true}, "legacy_history": JSON{"original": "preserved"}, "nodes": []FlowNode{}, "edges": []FlowEdge{}} if err := s.db.Save(r).Error; err != nil { t.Fatal(err) } original := compact(r.Data) if w := updateWorkbenchResource(s, ctx, r, resourceRequest{Name: "renamed", Version: 1, Data: r.Data}); w.Code != 200 { t.Fatal(w.Code, w.Body.String()) } r, _ = s.get(ctx, r.Kind, r.ID) if requireExecutable(r) == nil { t.Fatal("rename activated incomplete import") } forged := resourceData(r, r.Data) forged["legacy_activation"] = JSON{"status": "active"} if w := updateWorkbenchResource(s, ctx, r, resourceRequest{Name: r.Name, Version: r.Version, Data: forged, Configure: true}); w.Code != 400 { t.Fatal("forged activation accepted") } if w := updateWorkbenchResource(s, ctx, r, resourceRequest{Name: r.Name, Version: r.Version, Data: r.Data, Configure: true}); w.Code != 400 { t.Fatal("empty canvas activated") } configured := resourceData(r, FlowData{Nodes: []FlowNode{{ID: "n1", AgentID: "author"}}}) if w := updateWorkbenchResource(s, ctx, r, resourceRequest{Name: r.Name, Version: r.Version, Data: configured, Configure: true}); w.Code != 200 { t.Fatal(w.Code, w.Body.String()) } r, _ = s.get(ctx, r.Kind, r.ID) if err := requireExecutable(r); err != nil { t.Fatal(err) } if compact(r.Data["legacy_history"]) != `{"original":"preserved"}` { t.Fatal("history lost") } var revision Revision if err := s.db.First(&revision, "resource_id = ? AND version = 1", r.ID).Error; err != nil || compact(revision.Data) != original { t.Fatal("preedit import lost") } } func TestTaskConfigurationKeepsHistoricalOutput(t *testing.T) { d := TaskData{FlowID: "old", SourceText: "original", Segments: []Segment{{ID: "s", Content: "historical", AgentID: "a", Status: "accepted", Review: "segment review"}}, Review: "review", Composed: "final"} r := &Resource{Kind: "content-tasks", Data: asJSON(d), Status: "completed"} d.FlowID = "new" if err := applyManualUpdate(r, resourceRequest{Name: "task", Data: asJSON(d)}); err != nil { t.Fatal(err) } got, _ := decode[TaskData](r.Data) if len(got.Segments) != 1 || got.Segments[0].Content != "historical" || got.Segments[0].Status != "needs_review" || got.Review != "review" || got.Composed != "final" || r.Status != "needs_review" { t.Fatalf("output lost or stale approval retained: %+v", got) } } func TestTextTaskInheritsAgentKnowledgeAndIgnoresLegacyCanvasScope(t *testing.T) { s := sqliteWorkbench(t) s.agents = generationAgentStub{config: types.CustomAgentConfig{KnowledgeBases: []string{"agent-kb"}, KBSelectionMode: "selected"}} ctx := workbenchContext(1, "alice", types.TenantRoleContributor) f := seedResource(t, s, 1, "alice", "flows") f.Data = asJSON(FlowData{Nodes: []FlowNode{{ID: "n", AgentID: "canvas-agent"}}, KnowledgeIDs: []string{"canvas-file"}, KnowledgeFolders: map[string][]string{"kb": {"folder"}}}) if err := s.db.Save(f).Error; err != nil { t.Fatal(err) } r := &Resource{Kind: "content-tasks", Name: "task", Data: asJSON(TaskData{SourceText: "manual framework", FlowID: f.ID, AgentID: "injected", KBIDs: []string{"injected-kb"}})} if err := s.validate(ctx, r); err != nil { t.Fatal(err) } d, _ := decode[TaskData](r.Data) if d.AgentID != "canvas-agent" || len(d.KBIDs) != 1 || d.KBIDs[0] != "agent-kb" || len(d.KnowledgeIDs) != 0 || len(d.KnowledgeFolders) != 0 || len(d.KnowledgeTags) != 0 { t.Fatalf("task did not inherit: %+v", d) } if err := s.inheritFlow(ctx, &d, FlowData{Nodes: []FlowNode{{ID: "n", AgentID: "canvas-agent"}}}); err != nil || len(d.KBIDs) != 1 || d.KBIDs[0] != "agent-kb" { t.Fatalf("agent scope not inherited: %+v %v", d, err) } } func TestAIRecommendsTrainingDays(t *testing.T) { s, ctx, job, model := generationHarness(t) job.Action = "generate" model.output = `{"days":[{"day":1,"title":"first","hook":"next","nodes":[{"id":"n1","label":"foundation"}]},{"day":2,"title":"second","prerequisite":"first","hook":"apply","nodes":[{"id":"n2","label":"practice"}]}]}` r := &Resource{Kind: "training-maps", Data: asJSON(MapData{Topic: "topic", AgentID: "author"})} if err := s.executeMap(ctx, job, r, ActionRequest{}); err != nil { t.Fatal(err) } d, _ := decode[MapData](r.Data) if d.DaysCount != 2 || !d.DaysAuto || !strings.Contains(model.messages[1].Content, "推荐合理天数") { t.Fatal("AI day recommendation not persisted") } if validateDays(nil, 0) == nil { t.Fatal("zero recommended days accepted") } } func TestNodeRevisionOnlyReplacesRequestedSubtree(t *testing.T) { s, ctx, job, model := generationHarness(t) job.Action = "revise" model.output = `{"days":[{"day":1,"title":"unrequested change","hook":"unrequested","nodes":[{"id":"a","label":"changed","children":[{"id":"a1","label":"child"}]},{"id":"b","label":"unrequested sibling change"}]}]}` r := &Resource{Data: asJSON(MapData{AgentID: "author", DaysCount: 1, Days: []Day{{Day: 1, Title: "original title", Hook: "original hook", Nodes: []Node{{ID: "a", Label: "before"}, {ID: "b", Label: "sibling"}}}}})} if err := s.executeMap(ctx, job, r, ActionRequest{Day: 1, NodeID: "a", Advice: "fix"}); err != nil { t.Fatal(err) } d, _ := decode[MapData](r.Data) if d.Days[0].Title != "original title" || d.Days[0].Hook != "original hook" || d.Days[0].Nodes[1].Label != "sibling" || d.Days[0].Nodes[0].Label != "changed" { t.Fatalf("unrelated content changed: %+v", d.Days[0]) } if err := s.executeMap(ctx, job, r, ActionRequest{Day: 1, NodeID: "missing", Advice: "fix"}); !errors.Is(err, ErrInvalid) { t.Fatal("unknown node accepted") } } type imageModelStub struct { interfaces.ModelService model vlm.VLM id string } func (s *imageModelStub) GetVLMModel(_ context.Context, id string) (vlm.VLM, error) { s.id = id return s.model, nil } func (s *imageModelStub) ListModels(context.Context) ([]*types.Model, error) { return []*types.Model{{ID: "workspace-vision", Type: types.ModelTypeVLLM, Status: types.ModelStatusActive}}, nil } type imageVLMStub struct { vlm.VLM seen bool } func (v *imageVLMStub) Predict(_ context.Context, data [][]byte, prompt string) (string, error) { v.seen = len(data) == 1 && strings.Contains(prompt, "节点") return "# recognized framework", nil } func TestImageOCRUsesAgentModelAndWorkspaceFallback(t *testing.T) { image := "data:image/png;base64," + base64.StdEncoding.EncodeToString([]byte("\x89PNG\r\n\x1a\nimage")) for _, configured := range []string{"agent-vision", ""} { s := NewService(nil) v := &imageVLMStub{} m := &imageModelStub{model: v} s.models, s.agents = m, generationAgentStub{config: types.CustomAgentConfig{VLMModelID: configured}} text, err := s.recognizeSourceImage(workbenchContext(1, "alice", types.TenantRoleContributor), "agent", image) if err != nil || text != "# recognized framework" || !v.seen { t.Fatal(text, err) } want := configured if want == "" { want = "workspace-vision" } if m.id != want { t.Fatal("wrong VLM", m.id) } } for _, invalid := range []string{"https://external/image.png", "data:image/png;base64,not-base64", "data:image/png;base64," + base64.StdEncoding.EncodeToString([]byte("text"))} { if _, err := decodeSourceImage(invalid); err == nil { t.Fatal("invalid image accepted") } } }