xinmi/internal/workbench/structured_generation_test.go

package workbench

import (
	"context"
	"errors"
	"strings"
	"testing"

	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/types"
)

type repairChatStub struct {
	embeddedChat
	outputs      []string
	finishReason string
	err          error
	messages     [][]chat.Message
	options      []chat.ChatOptions
	attribution  []meteringAttribution
}

func (m *repairChatStub) GetModelID() string   { return "model" }
func (m *repairChatStub) GetModelName() string { return "repair-test" }

func (m *repairChatStub) Chat(ctx context.Context, messages []chat.Message, opts *chat.ChatOptions) (*types.ChatResponse, error) {
	index := len(m.messages)
	m.messages = append(m.messages, append([]chat.Message(nil), messages...))
	m.options = append(m.options, *opts)
	attribution, _ := ctx.Value(meteringContextKey{}).(meteringAttribution)
	m.attribution = append(m.attribution, attribution)
	if m.err != nil {
		return nil, m.err
	}
	output := m.outputs[min(index, len(m.outputs)-1)]
	return &types.ChatResponse{Content: output, FinishReason: m.finishReason, Usage: types.TokenUsage{PromptTokens: 80, CompletionTokens: 20, TotalTokens: 100}}, nil
}

const validMapResponse = `{"days":[{"day":1,"title":"基础","prerequisite":"","hook":"衔接进阶","nodes":[{"id":"n1","label":"核心概念","children":[]}]}]}`

func repairHarness(t *testing.T, outputs ...string) (*Service, context.Context, *Job, *Resource, *repairChatStub) {
	t.Helper()
	s, ctx, job, _ := generationHarness(t)
	model := &repairChatStub{outputs: outputs, finishReason: "stop"}
	s.models = generationModelStub{model: model}
	job.Action = "revise"
	r := &Resource{Status: "confirmed", Data: asJSON(MapData{AgentID: "author", Topic: "学习方法", DaysCount: 1, Days: []Day{{Day: 1, Title: "旧基础", Hook: "旧衔接", Nodes: []Node{{ID: "n1", Label: "原节点"}}, Framework: "已采用的框架", Confirmed: true}}})}
	return s, ctx, job, r, model
}

func TestStructuredOutputDiagnosticsNeverContainRawValues(t *testing.T) {
	for _, tc := range []struct {
		text, kind, field string
	}{
		{"```json\n{\"days\":[\"private-body-secret\",]}\n```", "json_syntax", ""},
		{`{"days":[{"day":"private-body-secret"}]}`, "json_type", "days.day"},
	} {
		_, err := parseJSON[mapDaysOutput](tc.text)
		var diagnostic *structuredOutputError
		if !errors.As(err, &diagnostic) || !errors.Is(err, ErrInvalid) {
			t.Fatalf("missing structured error: %v", err)
		}
		if diagnostic.Kind != tc.kind || diagnostic.Field != tc.field || diagnostic.Offset <= 0 || diagnostic.ResponseBytes != len(tc.text) {
			t.Fatalf("incorrect diagnostic: %+v", diagnostic)
		}
		if strings.Contains(err.Error()+compact(diagnostic.StructuredOutputDiagnostic()), "private-body-secret") {
			t.Fatal("private model output leaked into diagnostics")
		}
	}
	if safeJSONField("days.private-body-secret") != "unknown" {
		t.Fatal("unrecognized field name was exposed")
	}
}

func TestMapRepairsMalformedOutputOnceWithSameEvidenceAndMetering(t *testing.T) {
	for _, broken := range []string{
		`{"days":[`,
		`{"days":[{"day":"first"}]}`,
		`{"days":[{"day":1,"title":"基础","hook":"衔接","nodes":[{"id":"n1","label":"一"},{"id":"n1","label":"二"}]}]}`,
	} {
		s, ctx, job, r, model := repairHarness(t, broken, validMapResponse)
		data, _ := decode[MapData](r.Data)
		data.KBIDs = []string{"kb"}
		r.Data = asJSON(data)
		search := &budgetSearchStub{rows: []*types.SearchResult{{ID: "chunk", Content: "采用的课程依据"}}}
		s.sessions = search
		if err := s.executeMap(ctx, job, r, ActionRequest{Advice: "调整框架"}); err != nil {
			t.Fatal(err)
		}
		if len(model.messages) != 2 || search.calls != 1 {
			t.Fatalf("calls=%d retrievals=%d", len(model.messages), search.calls)
		}
		if model.messages[0][1].Content != model.messages[1][1].Content || !strings.Contains(model.messages[1][2].Content, compact(broken)) {
			t.Fatal("repair lost original evidence or malformed candidate")
		}
		for i := range model.options {
			if len(model.options[i].Format) == 0 || model.options[i].MaxCompletionTokens != model.options[0].MaxCompletionTokens || model.attribution[i].JobID != job.ID || model.attribution[i].AgentID != "author" {
				t.Fatal("repair bypassed JSON output mode, output cap or billing attribution")
			}
		}
		var saved Job
		if err := s.db.First(&saved, "id=?", job.ID).Error; err != nil {
			t.Fatal(err)
		}
		metadata, err := decode[struct {
			Attempt          int              `json:"attempt"`
			Usage            types.TokenUsage `json:"usage"`
			PreviousAttempts []struct {
				Usage      types.TokenUsage `json:"usage"`
				Structured JSON             `json:"structured_output"`
			} `json:"previous_attempts"`
		}](saved.Generation)
		if err != nil || metadata.Attempt != 2 || metadata.Usage.TotalTokens != 100 || len(metadata.PreviousAttempts) != 1 || metadata.PreviousAttempts[0].Usage.TotalTokens != 100 || metadata.PreviousAttempts[0].Structured["kind"] == nil {
			t.Fatalf("generation attempts/usage/diagnostics were not retained: %s", compact(saved.Generation))
		}
		out, _ := decode[MapData](r.Data)
		if out.Days[0].Title != "基础" || out.Days[0].Confirmed || out.Days[0].Framework != "" || r.Status != "draft" {
			t.Fatal("repaired output did not follow normal approval invalidation")
		}
	}
}

func TestMapRejectsSecondInvalidOutputAndPreservesAcceptedVersion(t *testing.T) {
	for _, output := range []string{`{"days":[`, `{"days":[{"day":1,"title":"基础","hook":"衔接","nodes":[{"id":"n1","label":"一"},{"id":"n1","label":"二"}]}]}`} {
		s, ctx, job, r, model := repairHarness(t, output)
		original := compact(r.Data)
		err := s.executeMap(ctx, job, r, ActionRequest{Advice: "调整框架"})
		if !errors.Is(err, ErrInvalid) || len(model.messages) != 2 || !strings.Contains(err.Error(), "已自动修复一次") {
			t.Fatalf("unexpected bounded retry result: %v, calls=%d", err, len(model.messages))
		}
		if compact(r.Data) != original || r.Status != "confirmed" {
			t.Fatal("failed repair changed the accepted resource")
		}
	}
}

func TestMapRepairSettlesEachModelCallExactlyOnce(t *testing.T) {
	s, ctx, job, r, model := repairHarness(t, `{"days":[`, validMapResponse)
	if err := s.db.AutoMigrate(&Balance{}, &BalanceHold{}, &Usage{}, &UsageQuota{}, &Transaction{}); err != nil {
		t.Fatal(err)
	}
	seedBalance(t, s, 1, "writer", 1_000_000)
	s.models = generationModelStub{model: s.WrapChat(model)}
	if err := s.executeMap(ctx, job, r, ActionRequest{}); err != nil {
		t.Fatal(err)
	}
	var calls []Usage
	if err := s.db.Where("job_id=?", job.ID).Find(&calls).Error; err != nil {
		t.Fatal(err)
	}
	if len(calls) != 2 || calls[0].CallID == calls[1].CallID {
		t.Fatalf("repair billing calls were lost or deduplicated: %+v", calls)
	}
	account := accountBalance(t, s, 1, "writer")
	if account.Balance != 1_000_000-200*pointsPerToken || account.Reserved != 0 {
		t.Fatalf("repair left an incorrect balance or reservation: %+v", account)
	}
	var debitCount int64
	if err := s.db.Model(&Transaction{}).Count(&debitCount).Error; err != nil || debitCount != 2 {
		t.Fatalf("incorrect debit count %d: %v", debitCount, err)
	}
}

func TestMapValidOutputSkipsRepairAndDailyFrameworkSupportsRepair(t *testing.T) {
	s, ctx, job, r, model := repairHarness(t, validMapResponse)
	if err := s.executeMap(ctx, job, r, ActionRequest{}); err != nil || len(model.messages) != 1 {
		t.Fatalf("valid output retried: %v", err)
	}
	s, ctx, job, r, model = repairHarness(t, `{"framework":["节点"]}`, `{"framework":"# 基础\n- 节点"}`)
	job.Action = "generate-day"
	if err := s.executeMap(ctx, job, r, ActionRequest{Day: 1}); err != nil || len(model.messages) != 2 {
		t.Fatalf("daily repair failed: %v", err)
	}
	out, _ := decode[MapData](r.Data)
	if out.Days[0].Framework != "# 基础\n- 节点" || out.Days[0].Confirmed {
		t.Fatal("daily repaired output was not saved for confirmation")
	}
}

func TestMapRepairHonorsContextBudgetAndDoesNotRetryProviderFailures(t *testing.T) {
	s, ctx, job, r, model := repairHarness(t, strings.Repeat("private-output ", 6000))
	s.agents = generationAgentStub{config: types.CustomAgentConfig{MaxCompletionTokens: 512}}
	s.models = generationModelStub{model: model, parameters: types.ModelParameters{ContextWindow: 4096}}
	original := compact(r.Data)
	err := s.executeMap(ctx, job, r, ActionRequest{})
	if !errors.Is(err, ErrInvalid) || !strings.Contains(err.Error(), "上下文容量不足") || len(model.messages) != 1 || compact(r.Data) != original || r.Status != "confirmed" {
		t.Fatalf("repair bypassed context budget or mutated accepted content: %v, calls=%d", err, len(model.messages))
	}
	for _, reason := range []string{"provider", "length"} {
		s, ctx, job, r, model = repairHarness(t, `{"days":[`)
		if reason == "provider" {
			model.err = context.DeadlineExceeded
		} else {
			model.finishReason = "length"
		}
		if err := s.executeMap(ctx, job, r, ActionRequest{}); err == nil || len(model.messages) != 1 {
			t.Fatalf("provider failure triggered automatic format repair: %v", err)
		}
	}
}

func TestTargetedMapRevisionUsesActualDayAndRejectsInvalidSubtree(t *testing.T) {
	s, ctx, job, r, model := repairHarness(t, `{"days":[{"day":2,"nodes":[{"id":"n2","label":"新节点","children":[{"id":"n2","label":"重复编号"}]}]}]}`, `{"days":[{"day":2,"nodes":[{"id":"n2","label":"新节点"}]}]}`)
	d, _ := decode[MapData](r.Data)
	d.DaysCount = 2
	d.Days = append(d.Days, Day{Day: 2, Title: "进阶", Prerequisite: "第一天", Hook: "总结", Nodes: []Node{{ID: "n2", Label: "旧节点"}}, Framework: "已确认进阶", Confirmed: true})
	r.Data = asJSON(d)
	if err := s.executeMap(ctx, job, r, ActionRequest{Day: 2, NodeID: "n2", Advice: "修正节点"}); err != nil {
		t.Fatal(err)
	}
	if len(model.messages) != 2 || !strings.Contains(model.messages[0][0].Content, `"day":2`) {
		t.Fatal("targeted output example still describes the wrong day")
	}
	out, _ := decode[MapData](r.Data)
	if compact(out.Days[0]) != compact(d.Days[0]) || out.Days[1].Title != d.Days[1].Title || out.Days[1].Nodes[0].Label != "新节点" || out.Days[1].Confirmed {
		t.Fatal("targeted repair changed an unrelated day or bypassed confirmation")
	}
}

func TestMapValidationRetainsNodeAndContinuityLimits(t *testing.T) {
	chain := func(depth int) []Node {
		var nodes []Node
		for i := depth; i > 0; i-- {
			nodes = []Node{{ID: strings.Repeat("n", i), Label: "节点", Children: nodes}}
		}
		return nodes
	}
	if err := validateNodes(chain(5), 0, map[string]bool{}); err != nil {
		t.Fatal(err)
	}
	for _, nodes := range [][]Node{chain(6), {{ID: "n", Label: strings.Repeat("字", 101)}}, {{ID: " ", Label: "节点"}}, {{ID: "n", Label: "节点", Children: []Node{{ID: "n", Label: "重复"}}}}} {
		if err := validateNodes(nodes, 0, map[string]bool{}); !errors.Is(err, ErrInvalid) {
			t.Fatalf("invalid nodes passed validation: %v", err)
		}
	}
	for _, days := range [][]Day{
		{{Day: 1, Title: "基础", Nodes: chain(1)}},
		{{Day: 2, Title: "基础", Hook: "衔接", Nodes: chain(1)}},
		{{Day: 1, Title: "基础", Hook: "衔接", Nodes: chain(1)}, {Day: 2, Title: "进阶", Hook: "结束", Nodes: chain(1)}},
		{{Day: 1, Title: "基础", Hook: "衔接", Nodes: chain(1)}, {Day: 2, Title: "基础", Prerequisite: "基础", Hook: "结束", Nodes: chain(1)}},
	} {
		if err := validateDays(days, len(days)); !errors.Is(err, ErrInvalid) {
			t.Fatalf("invalid daily continuity passed: %v", err)
		}
	}
}