xinmi/internal/workbench/generation_test.go

package workbench

import (
	"context"
	"strings"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

type generationAgentStub struct {
	interfaces.CustomAgentService
	config types.CustomAgentConfig
}

func (a generationAgentStub) GetAgentByID(ctx context.Context, id string) (*types.CustomAgent, error) {
	tenant, _ := types.TenantIDFromContext(ctx)
	config := a.config
	if config.ModelID == "" {
		config.ModelID = "model"
	}
	return &types.CustomAgent{ID: id, TenantID: tenant, Config: config}, nil
}

type generationModelStub struct {
	interfaces.ModelService
	model      chat.Chat
	parameters types.ModelParameters
}

func (m generationModelStub) GetModelByID(context.Context, string) (*types.Model, error) {
	return &types.Model{ID: "model", Parameters: m.parameters}, nil
}

func (m generationModelStub) GetChatModel(context.Context, string) (chat.Chat, error) {
	return m.model, nil
}

type embeddedChat interface{ chat.Chat }

type generationChatStub struct {
	embeddedChat
	messages []chat.Message
	output   string
	options  chat.ChatOptions
	usage    types.TokenUsage
	err      error
}

func (m *generationChatStub) Chat(_ context.Context, messages []chat.Message, opts *chat.ChatOptions) (*types.ChatResponse, error) {
	m.messages = messages
	m.options = *opts
	return &types.ChatResponse{Content: m.output, Usage: m.usage}, m.err
}

func generationHarness(t *testing.T) (*Service, context.Context, *Job, *generationChatStub) {
	t.Helper()
	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
	if err != nil {
		t.Fatal(err)
	}
	sqlDB, _ := db.DB()
	sqlDB.SetMaxOpenConns(1)
	t.Cleanup(func() { _ = sqlDB.Close() })
	if err = db.AutoMigrate(&Job{}, &Resource{}, &Consumption{}); err != nil {
		t.Fatal(err)
	}
	ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(1))
	ctx = context.WithValue(ctx, types.UserIDContextKey, "writer")
	ctx = context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleContributor)
	ctx = context.WithValue(ctx, leaseContextKey{}, jobLease{"job", "lease"})
	until := time.Now().Add(time.Minute)
	j := &Job{ID: "job", TenantID: 1, OwnerID: "writer", Status: "running", LeaseToken: "lease", LeaseUntil: &until}
	if err = db.Create(j).Error; err != nil {
		t.Fatal(err)
	}
	m := &generationChatStub{output: "AI 输出"}
	s := NewService(db)
	s.SetDependencies(generationModelStub{model: m}, generationAgentStub{}, nil)
	return s, ctx, j, m
}

func TestGenerateContinuesAcceptedContentWithoutSkippedOrFutureDrafts(t *testing.T) {
	s, ctx, job, model := generationHarness(t)
	job.Action = "generate"
	d := TaskData{AgentID: "author", Segments: []Segment{
		{ID: "a", Title: "基础", Status: "accepted", Content: "采用的论据"},
		{ID: "b", Title: "删除主题", Status: "skipped", Content: "跳过的秘密正文"},
		{ID: "c", Title: "当前主题", AgentID: "author", Status: "pending"},
		{ID: "d", Title: "后续主题", Status: "accepted", Content: "未来草稿不应出现"},
	}, Review: "旧报告", Composed: "旧成稿"}
	r := &Resource{Data: asJSON(d), Status: "completed"}
	if err := s.executeTask(ctx, job, r, ActionRequest{SegmentID: "c", AgentID: "replacement"}); err != nil {
		t.Fatal(err)
	}
	prompt := model.messages[1].Content
	for _, rejected := range []string{"跳过的秘密正文", "删除主题", "未来草稿不应出现"} {
		if strings.Contains(prompt, rejected) {
			t.Fatalf("excluded content leaked into prompt: %s", rejected)
		}
	}
	if !strings.Contains(prompt, "采用的论据") || !strings.Contains(prompt, "后续主题") {
		t.Fatal("missing accepted context or future outline")
	}
	out, _ := decode[TaskData](r.Data)
	if out.Segments[2].AgentID != "replacement" || out.Segments[2].Status != "generated" {
		t.Fatal("agent replacement not recorded")
	}
	if out.Segments[3].Status != "needs_review" || out.Review != "" || out.Composed != "" {
		t.Fatal("dependent approvals were not invalidated")
	}
}

func TestGenerationWaitsForPredecessorDecision(t *testing.T) {
	d := TaskData{Segments: []Segment{{ID: "a", Status: "generated", Content: "未确认"}, {ID: "b", Status: "pending"}}}
	if _, err := segmentContext(d, 1, false); err == nil {
		t.Fatal("unaccepted predecessor was allowed")
	}
}

func TestReviewCallsAIAndComposePreservesAcceptedOriginals(t *testing.T) {
	for _, action := range []string{"review", "compose"} {
		t.Run(action, func(t *testing.T) {
			s, ctx, job, model := generationHarness(t)
			job.Action = action
			input := TaskData{AgentID: "author", Segments: []Segment{{ID: "a", Status: "accepted", Content: "已确认正文"}, {ID: "b", Status: "skipped", Content: "不应恢复的正文"}}}
			accepted, _ := acceptedManuscript(input)
			input.Quality = &QualityReport{Passed: true, Summary: "通过", ManuscriptHash: manuscriptHash(accepted)}
			if action == "review" {
				model.output = `{"passed":true,"summary":"衔接清晰","issues":[]}`
			}
			r := &Resource{Data: asJSON(input)}
			if err := s.executeTask(ctx, job, r, ActionRequest{}); err != nil {
				t.Fatal(err)
			}
			if action == "review" && (len(model.messages) != 2 || strings.Contains(model.messages[1].Content, "不应恢复的正文")) {
				t.Fatal("incorrect AI manuscript")
			}
			d, _ := decode[TaskData](r.Data)
			if action == "review" && (d.Quality == nil || !d.Quality.Passed || !strings.Contains(d.Review, "衔接清晰")) {
				t.Fatal("AI review not saved")
			}
			if action == "compose" && (d.Composed != "已确认正文" || r.Status != "completed" || len(model.messages) != 0) {
				t.Fatal("composition must preserve original text without calling AI")
			}
		})
	}
}

func TestNormalizeSegmentsRejectsDuplicateAndUnauthorizedAgent(t *testing.T) {
	f := FlowData{Nodes: []FlowNode{{ID: "n", AgentID: "allowed"}}}
	for _, segments := range [][]Segment{
		{{ID: "s", Title: "一", AgentID: "allowed"}, {ID: "s", Title: "二", AgentID: "allowed"}},
		{{ID: "a", Title: "重复", AgentID: "allowed"}, {ID: "b", Title: "重复", AgentID: "allowed"}},
		{{ID: "a", Title: "标题", AgentID: "another-user-agent"}},
	} {
		if normalizeSegments(segments, "allowed", f, nil) == nil {
			t.Fatal("invalid model output accepted")
		}
	}
	segments := []Segment{{Title: "正确标题", Status: "accepted", Content: "注入的正文"}}
	if err := normalizeSegments(segments, "allowed", f, []string{"source"}); err != nil {
		t.Fatal(err)
	}
	if segments[0].Status != "pending" || segments[0].Content != "" || segments[0].AgentID != "allowed" {
		t.Fatal("unsafe generated state retained")
	}
}

func TestDayRevisionInvalidatesDependentFrameworks(t *testing.T) {
	s, ctx, job, model := generationHarness(t)
	job.Action = "revise"
	model.output = `{"days":[{"day":1,"title":"新基础","hook":"下一步","nodes":[{"id":"n1","label":"新概念"}],"framework":"模型伪造框架","confirmed":true}]}`
	d := MapData{AgentID: "author", DaysCount: 2, KnowledgeFolders: map[string][]string{"kb": {"课程/基础"}}, Days: []Day{
		{Day: 1, Title: "基础", Hook: "下一步", Nodes: []Node{{ID: "n1", Label: "概念"}}, Framework: "旧框架", Confirmed: true},
		{Day: 2, Title: "进阶", Prerequisite: "基础", Hook: "总结", Nodes: []Node{{ID: "n2", Label: "进阶"}}, Framework: "旧下游框架", Confirmed: true},
	}}
	r := &Resource{Data: asJSON(d), Status: "confirmed"}
	if err := s.executeMap(ctx, job, r, ActionRequest{Day: 1, Advice: "修改基础"}); err != nil {
		t.Fatal(err)
	}
	out, _ := decode[MapData](r.Data)
	if len(out.KnowledgeFolders["kb"]) != 1 || out.KnowledgeFolders["kb"][0] != "课程/基础" {
		t.Fatal("selected knowledge folders disappeared during generation")
	}
	for _, day := range out.Days {
		if day.Confirmed || day.Framework != "" {
			t.Fatal("stale approval survived upstream revision")
		}
	}
	if r.Status != "draft" {
		t.Fatal("map remained confirmed")
	}
}