WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/workbench/generation_budget_test.go

@@ -0,0 +1,173 @@
+package workbench
+
+import (
+	"context"
+	"errors"
+	"reflect"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/agent/token"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+)
+
+func TestCompletionRespectsAgentAndModelOutputLimits(t *testing.T) {
+	for _, tc := range []struct {
+		name               string
+		agent, model, want int
+	}{
+		{"agent lower", 512, 8192, 512},
+		{"model lower", 8192, 1024, 1024},
+		{"default", 0, 0, types.DefaultQuickAnswerMaxCompletionTokens},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			s, ctx, _, model := generationHarness(t)
+			s.agents = generationAgentStub{config: types.CustomAgentConfig{MaxCompletionTokens: tc.agent}}
+			s.models = generationModelStub{model: model, parameters: types.ModelParameters{ContextWindow: 16000, MaxOutputTokens: tc.model}}
+			if _, err := s.complete(ctx, "agent", "system", "prompt"); err != nil {
+				t.Fatal(err)
+			}
+			if model.options.MaxCompletionTokens != tc.want {
+				t.Fatalf("output limit = %d, want %d", model.options.MaxCompletionTokens, tc.want)
+			}
+		})
+	}
+}
+
+func TestOversizedAcceptedContinuityFailsBeforeChatAndPreservesResource(t *testing.T) {
+	for _, action := range []string{"generate", "review"} {
+		t.Run(action, func(t *testing.T) {
+			s, ctx, job, model := generationHarness(t)
+			s.models = generationModelStub{model: model, parameters: types.ModelParameters{ContextWindow: 4096, MaxOutputTokens: 512}}
+			job.Action = action
+			d := TaskData{AgentID: "agent", Segments: []Segment{
+				{ID: "a", Title: "基础", Status: "accepted", Content: strings.Repeat("已确认的中文正文需要完整保留。", 1500)},
+				{ID: "b", Title: "进阶", AgentID: "agent", Status: "accepted", Content: "后文"},
+			}}
+			accepted, _ := acceptedManuscript(d)
+			d.Quality = &QualityReport{Passed: true, Summary: "通过", ManuscriptHash: manuscriptHash(accepted)}
+			r := &Resource{Data: asJSON(d)}
+			original := compact(r.Data)
+			err := s.executeTask(ctx, job, r, ActionRequest{SegmentID: "b"})
+			if !errors.Is(err, ErrInvalid) || !strings.Contains(err.Error(), "上下文容量不足") {
+				t.Fatalf("expected actionable context error, got %v", err)
+			}
+			if len(model.messages) != 0 || compact(r.Data) != original {
+				t.Fatal("oversized prompt reached model or changed accepted content")
+			}
+		})
+	}
+}
+
+type budgetSearchStub struct {
+	interfaces.SessionService
+	rows  []*types.SearchResult
+	calls int
+}
+
+func (s *budgetSearchStub) SearchKnowledge(context.Context, []string, []string, []types.TagScope, string) ([]*types.SearchResult, error) {
+	s.calls++
+	return s.rows, nil
+}
+
+func TestEvidenceBudgetIncludesSourceLabelsAndKeepsRelevanceOrder(t *testing.T) {
+	s := NewService(nil)
+	s.sessions = &budgetSearchStub{rows: []*types.SearchResult{
+		nil,
+		{ID: "empty", Content: "  "},
+		{ID: "first", KnowledgeFilename: "中文文件名", Content: "最高相关的事实。"},
+		{ID: "first", Content: "重复内容不可引用"},
+		{ID: "second", KnowledgeFilename: "另一份资料", Content: strings.Repeat("第二相关片段包含很多汉字。", 1000)},
+		{ID: "third", Content: "不应挤掉更相关片段"},
+	}}
+	const budget = 250
+	evidence, ids, err := s.evidence(context.Background(), []string{"kb"}, nil, "query", budget)
+	if err != nil {
+		t.Fatal(err)
+	}
+	estimator, _ := token.NewEstimator()
+	if n := estimator.EstimateString(evidence); n > budget {
+		t.Fatalf("evidence uses %d tokens, limit %d", n, budget)
+	}
+	if !reflect.DeepEqual(ids, []string{"first", "second"}) || !strings.Contains(evidence, "仅引用部分内容") || strings.Contains(evidence, "重复内容不可引用") {
+		t.Fatalf("source selection = %v, text = %s", ids, evidence)
+	}
+}
+
+func TestEvidenceFitsRemainingContextWithoutTruncatingAcceptedText(t *testing.T) {
+	s, ctx, _, model := generationHarness(t)
+	s.agents = generationAgentStub{config: types.CustomAgentConfig{MaxCompletionTokens: 512, SystemPrompt: strings.Repeat("表达风格", 30)}}
+	s.models = generationModelStub{model: model, parameters: types.ModelParameters{ContextWindow: 4096}}
+	s.sessions = &budgetSearchStub{rows: []*types.SearchResult{{ID: "chunk", Content: strings.Repeat("资料的正文内容。", 5000)}}}
+	prompt := "已采用正文：" + strings.Repeat("前序完整内容。", 160)
+	if _, _, err := s.completeWithEvidence(ctx, "agent", "system", prompt, []string{"kb"}, nil, "query"); err != nil {
+		t.Fatal(err)
+	}
+	if !strings.Contains(model.messages[1].Content, prompt) {
+		t.Fatal("accepted continuity was truncated")
+	}
+	estimator, _ := token.NewEstimator()
+	if n := estimator.EstimateMessages(model.messages); n+512 > 4096-4096*15/100 {
+		t.Fatalf("full request exceeded context safety budget: %d", n)
+	}
+}
+
+func TestEvidenceRefusesFullContextBeforeRetrieval(t *testing.T) {
+	s := NewService(nil)
+	search := &budgetSearchStub{}
+	s.sessions = search
+	if _, _, err := s.evidence(context.Background(), []string{"kb"}, nil, "query", 0); !errors.Is(err, ErrInvalid) || search.calls != 0 {
+		t.Fatalf("full context should fail before retrieval: %v, calls=%d", err, search.calls)
+	}
+}
+
+func TestCompletionRecordsProviderUsageOnOwnedJob(t *testing.T) {
+	s, ctx, job, model := generationHarness(t)
+	model.usage = types.TokenUsage{PromptTokens: 80, CompletionTokens: 20, TotalTokens: 100}
+	if _, err := s.complete(ctx, "agent", "system", "prompt"); err != nil {
+		t.Fatal(err)
+	}
+	var saved Job
+	if err := s.db.First(&saved, "id = ?", job.ID).Error; err != nil {
+		t.Fatal(err)
+	}
+	metadata, err := decode[struct {
+		AgentID   string           `json:"agent_id"`
+		ModelID   string           `json:"model_id"`
+		Estimated int              `json:"estimated_input_tokens"`
+		Reported  bool             `json:"provider_usage_reported"`
+		Usage     types.TokenUsage `json:"usage"`
+	}](saved.Generation)
+	if err != nil || metadata.AgentID != "agent" || metadata.ModelID != "model" || metadata.Estimated <= 0 || !metadata.Reported || metadata.Usage.TotalTokens != 100 {
+		t.Fatalf("usage metadata missing: %+v, %v", saved.Generation, err)
+	}
+}
+
+func TestCompletionRequiresValidOwnedLeaseBeforeCallingModel(t *testing.T) {
+	for _, column := range []string{"owner_id", "lease_until"} {
+		t.Run(column, func(t *testing.T) {
+			s, ctx, job, model := generationHarness(t)
+			var value any = "someone-else"
+			if column == "lease_until" {
+				value = time.Now().Add(-time.Minute)
+			}
+			if err := s.db.Model(&Job{}).Where("id = ?", job.ID).Update(column, value).Error; err != nil {
+				t.Fatal(err)
+			}
+			if _, err := s.complete(ctx, "agent", "system", "prompt"); !errors.Is(err, ErrConflict) || len(model.messages) > 0 {
+				t.Fatalf("invalid lease reached provider: %v", err)
+			}
+		})
+	}
+}
+
+func TestProviderContextOverflowProducesActionableErrorWithoutRawDetails(t *testing.T) {
+	s, ctx, _, model := generationHarness(t)
+	model.err = errors.New("context_length_exceeded provider-secret-marker")
+	_, err := s.complete(ctx, "agent", "system", "prompt")
+	if !errors.Is(err, ErrInvalid) || !strings.Contains(publicJobError(err), "上下文") || strings.Contains(publicJobError(err), "provider-secret-marker") {
+		t.Fatalf("unexpected public error: %v", err)
+	}
+}