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

internal/application/service/extract_graph_budget_test.go

@@ -0,0 +1,186 @@
+package service
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"path/filepath"
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/Tencent/WeKnora/internal/models/chat"
+	"github.com/Tencent/WeKnora/internal/models/requestbudget"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/utils"
+	"github.com/hibiken/asynq"
+	"github.com/stretchr/testify/require"
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+	"gorm.io/gorm/logger"
+)
+
+const graphBudgetSource = "Trust supports learning."
+
+func graphBudgetPayload() types.ExtractChunkPayload {
+	return types.ExtractChunkPayload{TenantID: 7, ChunkID: "chunk", ModelID: "model",
+		ExpectedSource: &types.GraphSourceIdentity{KnowledgeBaseID: "kb", KnowledgeID: "doc", ContentSHA256: fmt.Sprintf("%x", sha256.Sum256([]byte(graphBudgetSource)))}}
+}
+
+func openGraphBudgetDB(t *testing.T, path string) *gorm.DB {
+	t.Helper()
+	db, err := gorm.Open(sqlite.Open(path), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
+	require.NoError(t, err)
+	sqlDB, err := db.DB()
+	require.NoError(t, err)
+	sqlDB.SetMaxOpenConns(1)
+	t.Cleanup(func() { _ = sqlDB.Close() })
+	require.NoError(t, db.Exec(`CREATE TABLE IF NOT EXISTS native_graph_backfill_failures (
+		task_id TEXT PRIMARY KEY, tenant_id INTEGER, knowledge_base_id TEXT, chunk_id TEXT,
+		knowledge_id TEXT, content_hash TEXT, state TEXT, manual_retries INTEGER, reason TEXT, updated_at DATETIME)`).Error)
+	return db
+}
+
+func seedGraphBudget(t *testing.T, db *gorm.DB) {
+	t.Helper()
+	p := graphBudgetPayload()
+	require.NoError(t, db.Exec(`INSERT INTO native_graph_backfill_failures VALUES(?,?,?,?,?,?,'dispatch_requested',0,'',CURRENT_TIMESTAMP)`,
+		"task", p.TenantID, p.ExpectedSource.KnowledgeBaseID, p.ChunkID, p.ExpectedSource.KnowledgeID, p.ExpectedSource.ContentSHA256).Error)
+}
+
+func graphBudgetDB(t *testing.T) *gorm.DB {
+	t.Helper()
+	db := openGraphBudgetDB(t, filepath.Join(t.TempDir(), "budget.sqlite"))
+	seedGraphBudget(t, db)
+	return db
+}
+
+func TestGraphExecutionClaimSurvivesRestartAndConcurrentDelivery(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "budget.sqlite")
+	db := openGraphBudgetDB(t, path)
+	seedGraphBudget(t, db)
+	svc := &ChunkExtractService{executionDB: db}
+	var successes atomic.Int32
+	var wg sync.WaitGroup
+	for i := 0; i < 8; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			if svc.claimGraphExecution(context.Background(), graphBudgetPayload(), "task") == nil {
+				successes.Add(1)
+			}
+		}()
+	}
+	wg.Wait()
+	require.Equal(t, int32(1), successes.Load())
+	sqlDB, _ := db.DB()
+	require.NoError(t, sqlDB.Close())
+	svc = &ChunkExtractService{executionDB: openGraphBudgetDB(t, path)}
+	require.ErrorIs(t, svc.claimGraphExecution(context.Background(), graphBudgetPayload(), "task"), requestbudget.ErrStopped)
+}
+
+func TestGraphExecutionClaimRejectsChangedOrFailedReceipt(t *testing.T) {
+	for _, changed := range []string{"task", "tenant", "base", "chunk", "document", "source", "failed", "retry", "reason", "missing db"} {
+		t.Run(changed, func(t *testing.T) {
+			db := graphBudgetDB(t)
+			p, task := graphBudgetPayload(), "task"
+			svc := &ChunkExtractService{executionDB: db}
+			switch changed {
+			case "task":
+				task = "other"
+			case "tenant":
+				p.TenantID++
+			case "base":
+				p.ExpectedSource.KnowledgeBaseID = "other"
+			case "chunk":
+				p.ChunkID = "other"
+			case "document":
+				p.ExpectedSource.KnowledgeID = "other"
+			case "source":
+				p.ExpectedSource.ContentSHA256 = "other"
+			case "failed":
+				require.NoError(t, db.Exec("UPDATE native_graph_backfill_failures SET state='failed'").Error)
+			case "retry":
+				require.NoError(t, db.Exec("UPDATE native_graph_backfill_failures SET manual_retries=1").Error)
+			case "reason":
+				require.NoError(t, db.Exec("UPDATE native_graph_backfill_failures SET reason='semantic_review_rejected'").Error)
+			case "missing db":
+				svc.executionDB = nil
+			}
+			require.ErrorIs(t, svc.claimGraphExecution(context.Background(), p, task), requestbudget.ErrStopped)
+		})
+	}
+}
+
+func TestGraphExecutionCorrectionSharesActualSendBudget(t *testing.T) {
+	t.Setenv("SSRF_WHITELIST", "127.0.0.1")
+	utils.ResetSSRFWhitelistForTest()
+	t.Cleanup(utils.ResetSSRFWhitelistForTest)
+	for _, tc := range []struct {
+		name          string
+		everyRedirect bool
+		wrongModel    bool
+	}{
+		{name: "three sends complete correction"},
+		{name: "fourth send refuses graph write", everyRedirect: true},
+		{name: "wrong concrete model sends nothing", wrongModel: true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			var sends, calls atomic.Int32
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				sends.Add(1)
+				if r.URL.Path != "/result" {
+					call := calls.Add(1)
+					if tc.everyRedirect || call == 2 {
+						http.Redirect(w, r, "/result", http.StatusTemporaryRedirect)
+						return
+					}
+				}
+				content := `[]`
+				if calls.Load() == 2 {
+					content = `[{"entity1":"Trust","entity2":"learning","relation":"supports","evidence":"Trust supports learning."}]`
+				}
+				w.Header().Set("Content-Type", "application/json")
+				_ = json.NewEncoder(w).Encode(map[string]any{"choices": []any{map[string]any{"message": map[string]any{"role": "assistant", "content": content}, "finish_reason": "stop"}}})
+			}))
+			defer server.Close()
+			modelID := "model"
+			if tc.wrongModel {
+				modelID = "other"
+			}
+			model, err := chat.NewChat(&chat.ChatConfig{Source: types.ModelSourceRemote, BaseURL: server.URL, ModelID: modelID, ModelName: "gpt-4o", APIKey: "local-only", Provider: "openai"}, nil)
+			require.NoError(t, err)
+			db := graphBudgetDB(t)
+			writer := &extractionWriter{}
+			svc := &ChunkExtractService{template: &types.PromptTemplateStructured{Description: "Extract graph"}, modelService: extractionModels{model: model},
+				knowledgeBaseRepo: extractionKB{}, knowledgeRepo: &extractionKnowledge{}, executionDB: db,
+				chunkRepo: extractionChunks{chunk: &types.Chunk{ID: "chunk", TenantID: 7, KnowledgeBaseID: "kb", KnowledgeID: "doc", Content: graphBudgetSource, IsEnabled: true, ChunkType: types.ChunkTypeText}}, graphEngine: writer}
+			payload, _ := json.Marshal(graphBudgetPayload())
+			task := asynq.NewTask(types.TypeChunkExtract, payload)
+			err = svc.handle(context.Background(), task, "task")
+			if tc.wrongModel {
+				require.Error(t, err)
+				require.Zero(t, sends.Load())
+				require.Nil(t, writer.saved)
+			} else if tc.everyRedirect {
+				require.Error(t, err)
+				require.Equal(t, int32(3), sends.Load())
+				require.Nil(t, writer.saved)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, int32(3), sends.Load())
+				require.True(t, writer.complete)
+			}
+			if !writer.complete {
+				before := sends.Load()
+				err = svc.handle(context.Background(), task, "task")
+				require.True(t, errors.Is(err, asynq.SkipRetry))
+				require.Equal(t, before, sends.Load(), "redelivery must not replenish sends")
+			}
+		})
+	}
+}