xinmi/internal/workbench/configuration_test.go

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