xinmi/internal/workbench/generation_budget_test.go

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