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

internal/application/service/wiki_failure_capture_test.go

@@ -0,0 +1,339 @@
+package service
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"os"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/agent"
+	"github.com/Tencent/WeKnora/internal/models/chat"
+	"github.com/Tencent/WeKnora/internal/models/requestbudget"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/stretchr/testify/require"
+)
+
+func wikiFailureFixture(t *testing.T) (context.Context, wikiFailureCaptureScope, map[string]string, string, error) {
+	t.Helper()
+	directory, err := filepath.EvalSymlinks(t.TempDir())
+	require.NoError(t, err)
+	require.NoError(t, os.Chmod(directory, 0700))
+	input := map[string]string{"Content": "Only an original name appears here.", "CustomInstructions": "UNRELATED_PRIVATE_INSTRUCTION"}
+	raw := "{\"entities\":[{\"name\":\"Unsupported name\",\"slug\":\"entity/unsupported\"}],\"concepts\":[]}\n"
+	result, validation := parseWikiCandidateOutput(raw)
+	require.Nil(t, validation)
+	original := validateWikiFreshSourceIdentities(result, input["Content"])
+	require.Error(t, original)
+	scope := wikiFailureCaptureScope{Directory: directory, TenantID: 71, KnowledgeBaseID: "kb", KnowledgeID: "doc", PendingOpID: 31,
+		Attempt: 2, ModelID: "model", Phase: "wiki_candidate_slug", SourceHash: wikiCheckpointHash([]string{input["Content"]}),
+		Release: "reviewed-release", APIImage: "reviewed-image", ExpiresAt: time.Now().Add(time.Hour)}
+	guard := &wikiRecoveryGuard{row: types.WikiRecoveryBudget{TenantID: scope.TenantID, KnowledgeBaseID: scope.KnowledgeBaseID,
+		KnowledgeID: scope.KnowledgeID, PendingOpID: scope.PendingOpID, AttemptNumber: scope.Attempt, ModelID: scope.ModelID,
+		ReviewedRelease: scope.Release, ReviewedAPIImage: scope.APIImage, PausedPayload: "UNRELATED_PRIVATE_PAYLOAD"}}
+	ctx := requestbudget.WithGuard(context.WithValue(context.Background(), types.TenantIDContextKey, scope.TenantID), guard)
+	return ctx, scope, input, raw, original
+}
+
+func setWikiFailureScope(t *testing.T, scope wikiFailureCaptureScope) {
+	t.Helper()
+	encoded, err := json.Marshal(scope)
+	require.NoError(t, err)
+	t.Setenv(wikiFailureCaptureEnv, string(encoded))
+}
+
+func TestWikiFailureCaptureRejectsServedDirectory(t *testing.T) {
+	ctx, scope, input, raw, original := wikiFailureFixture(t)
+	t.Setenv("LOCAL_STORAGE_BASE_DIR", scope.Directory)
+	setWikiFailureScope(t, scope)
+	require.Equal(t, "private_directory_is_served", retainWikiIdentityFailure(ctx, "model", agent.WikiCandidateSlugPrompt, input, "", raw, original))
+	files, err := os.ReadDir(scope.Directory)
+	require.NoError(t, err)
+	require.Empty(t, files)
+	require.True(t, wikiFailureDirectoryIsServed(filepath.Join(scope.Directory, "nested")))
+	require.False(t, wikiFailureDirectoryIsServed(scope.Directory+"-private"))
+	for _, directory := range []string{"/data/uploads/failure", "/usr/share/nginx/html/failure"} {
+		require.True(t, wikiFailureDirectoryIsServed(directory))
+	}
+	link := filepath.Join(t.TempDir(), "served-link")
+	require.NoError(t, os.Symlink(scope.Directory, link))
+	t.Setenv("LOCAL_STORAGE_BASE_DIR", link)
+	require.True(t, wikiFailureDirectoryIsServed(scope.Directory))
+}
+
+func readWikiFailureArtifact(t *testing.T, scope wikiFailureCaptureScope) ([]byte, wikiIdentityFailureArtifact) {
+	t.Helper()
+	files, err := filepath.Glob(filepath.Join(scope.Directory, "*", "failure.json"))
+	require.NoError(t, err)
+	require.Len(t, files, 1)
+	info, err := os.Stat(files[0])
+	require.NoError(t, err)
+	require.Equal(t, os.FileMode(0600), info.Mode().Perm())
+	info, err = os.Stat(filepath.Dir(files[0]))
+	require.NoError(t, err)
+	require.Equal(t, os.FileMode(0700), info.Mode().Perm())
+	encoded, err := os.ReadFile(files[0])
+	require.NoError(t, err)
+	var artifact wikiIdentityFailureArtifact
+	require.NoError(t, json.Unmarshal(encoded, &artifact))
+	return encoded, artifact
+}
+
+func TestWikiFailureCaptureDisabledAndWrongScope(t *testing.T) {
+	ctx, scope, input, raw, original := wikiFailureFixture(t)
+	t.Setenv(wikiFailureCaptureEnv, "")
+	require.Empty(t, retainWikiIdentityFailure(ctx, scope.ModelID, agent.WikiCandidateSlugPrompt, input, "", raw, original))
+	for name, change := range map[string]func(*wikiFailureCaptureScope){
+		"tenant":    func(s *wikiFailureCaptureScope) { s.TenantID++ },
+		"kb":        func(s *wikiFailureCaptureScope) { s.KnowledgeBaseID = "other" },
+		"document":  func(s *wikiFailureCaptureScope) { s.KnowledgeID = "other" },
+		"operation": func(s *wikiFailureCaptureScope) { s.PendingOpID++ },
+		"attempt":   func(s *wikiFailureCaptureScope) { s.Attempt++ },
+		"model":     func(s *wikiFailureCaptureScope) { s.ModelID = "other" },
+		"source":    func(s *wikiFailureCaptureScope) { s.SourceHash = strings.Repeat("a", 64) },
+		"release":   func(s *wikiFailureCaptureScope) { s.Release = "other" },
+		"image":     func(s *wikiFailureCaptureScope) { s.APIImage = "other" },
+	} {
+		t.Run(name, func(t *testing.T) {
+			other := scope
+			change(&other)
+			setWikiFailureScope(t, other)
+			require.Empty(t, retainWikiIdentityFailure(ctx, scope.ModelID, agent.WikiCandidateSlugPrompt, input, "", raw, original))
+		})
+	}
+	setWikiFailureScope(t, scope)
+	require.Empty(t, retainWikiIdentityFailure(context.Background(), scope.ModelID, agent.WikiCandidateSlugPrompt, input, "", raw, original))
+	require.Empty(t, retainWikiIdentityFailure(ctx, scope.ModelID, agent.WikiCandidateSlugPrompt, input, "", raw, errors.New("other error")))
+	require.Empty(t, retainWikiIdentityFailure(ctx, scope.ModelID, "other prompt", input, "", raw, original))
+	files, err := os.ReadDir(scope.Directory)
+	require.NoError(t, err)
+	require.Empty(t, files)
+}
+
+func TestWikiFailureCaptureExactPrivateImmutableReplay(t *testing.T) {
+	ctx, scope, input, raw, original := wikiFailureFixture(t)
+	setWikiFailureScope(t, scope)
+	require.Equal(t, "complete", retainWikiIdentityFailure(ctx, scope.ModelID, agent.WikiCandidateSlugPrompt, input, "repair", raw, original))
+	encoded, artifact := readWikiFailureArtifact(t, scope)
+	require.Equal(t, raw, artifact.Response)
+	require.Equal(t, input["Content"], artifact.Source)
+	require.Equal(t, wikiFailureTextHash(raw), artifact.ResponseHash)
+	require.Equal(t, wikiFailureTextHash(input["Content"]), artifact.SourceHash)
+	require.Equal(t, wikiCheckpointHash(input), artifact.InputHash)
+	require.Equal(t, wikiFailureTextHash("repair"), artifact.RepairHash)
+	require.NotContains(t, string(encoded), "UNRELATED_PRIVATE")
+	require.NotContains(t, string(encoded), scope.Directory)
+	replayed, validation := parseWikiCandidateOutput(artifact.Response)
+	require.Nil(t, validation)
+	require.EqualError(t, validateWikiFreshSourceIdentities(replayed, artifact.Source), original.Error())
+	require.Equal(t, "already_claimed", retainWikiIdentityFailure(ctx, scope.ModelID, agent.WikiCandidateSlugPrompt, input, "", "changed", original))
+	again, _ := readWikiFailureArtifact(t, scope)
+	require.Equal(t, encoded, again)
+}
+
+func TestWikiFailureCaptureBoundsPermissionsAndConcurrency(t *testing.T) {
+	t.Run("oversized response keeps metadata only", func(t *testing.T) {
+		ctx, scope, input, _, original := wikiFailureFixture(t)
+		setWikiFailureScope(t, scope)
+		raw := strings.Repeat("x", 1024*1024+1)
+		require.Equal(t, "size_limit", retainWikiIdentityFailure(ctx, scope.ModelID, agent.WikiCandidateSlugPrompt, input, "", raw, original))
+		_, artifact := readWikiFailureArtifact(t, scope)
+		require.Empty(t, artifact.Response)
+		require.Empty(t, artifact.Source)
+		require.Equal(t, len(raw), artifact.ResponseBytes)
+		require.Equal(t, wikiFailureTextHash(raw), artifact.ResponseHash)
+	})
+	t.Run("encoded size counts escaping", func(t *testing.T) {
+		ctx, scope, input, _, original := wikiFailureFixture(t)
+		setWikiFailureScope(t, scope)
+		raw := strings.Repeat("\x00", 512*1024)
+		require.Equal(t, "size_limit", retainWikiIdentityFailure(ctx, scope.ModelID, agent.WikiCandidateSlugPrompt, input, "", raw, original))
+		_, artifact := readWikiFailureArtifact(t, scope)
+		require.Empty(t, artifact.Response)
+	})
+	for _, mode := range []string{"permissions", "missing", "symlink"} {
+		t.Run(mode, func(t *testing.T) {
+			ctx, scope, input, raw, original := wikiFailureFixture(t)
+			before := original.Error()
+			want := "private_directory_unavailable"
+			switch mode {
+			case "permissions":
+				require.NoError(t, os.Chmod(scope.Directory, 0755))
+				want = "private_directory_permissions"
+			case "missing":
+				scope.Directory = filepath.Join(scope.Directory, "missing")
+			case "symlink":
+				link := filepath.Join(scope.Directory, "link")
+				require.NoError(t, os.Symlink(scope.Directory, link))
+				scope.Directory = link
+			}
+			setWikiFailureScope(t, scope)
+			require.Equal(t, want, retainWikiIdentityFailure(ctx, scope.ModelID, agent.WikiCandidateSlugPrompt, input, "", raw, original))
+			require.EqualError(t, original, before)
+		})
+	}
+	t.Run("one durable claim across concurrent calls", func(t *testing.T) {
+		ctx, scope, input, raw, original := wikiFailureFixture(t)
+		setWikiFailureScope(t, scope)
+		var wg sync.WaitGroup
+		outcomes := make(chan string, 8)
+		for range 8 {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				outcomes <- retainWikiIdentityFailure(ctx, scope.ModelID, agent.WikiCandidateSlugPrompt, input, "", raw, original)
+			}()
+		}
+		wg.Wait()
+		close(outcomes)
+		counts := map[string]int{}
+		for outcome := range outcomes {
+			counts[outcome]++
+		}
+		require.Equal(t, map[string]int{"complete": 1, "already_claimed": 7}, counts)
+		readWikiFailureArtifact(t, scope)
+	})
+}
+
+func TestWikiFailureCaptureDoesNotAddRequestsAfterIdentityPatchFailure(t *testing.T) {
+	t.Setenv(wikiFailureCaptureEnv, "invalid")
+	calls := 0
+	model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
+		calls++
+		return &types.ChatResponse{Content: `{"entities":[{"name":"unsupported","slug":"entity/unsupported"}],"concepts":[]}`, FinishReason: "stop"}, nil
+	}}
+	_, err := (&wikiIngestService{}).generateWikiFreshExtraction(context.Background(), model, agent.WikiCandidateSlugPrompt,
+		map[string]string{"Content": strings.Repeat("Original source.\n", 300)})
+	var identity *wikiSourceIdentityError
+	require.ErrorAs(t, err, &identity)
+	require.Equal(t, "wiki_identity_name_not_in_source", identity.Code)
+	require.Equal(t, 2, calls)
+}
+
+func TestWikiFailureCaptureInvalidConfigurations(t *testing.T) {
+	ctx, scope, input, raw, original := wikiFailureFixture(t)
+	for name, change := range map[string]func(*wikiFailureCaptureScope){
+		"expired":           func(s *wikiFailureCaptureScope) { s.ExpiresAt = time.Now().Add(-time.Second) },
+		"missing directory": func(s *wikiFailureCaptureScope) { s.Directory = "" },
+		"missing release":   func(s *wikiFailureCaptureScope) { s.Release = "" },
+		"missing source":    func(s *wikiFailureCaptureScope) { s.SourceHash = "" },
+		"wrong phase":       func(s *wikiFailureCaptureScope) { s.Phase = "wiki_page" },
+	} {
+		t.Run(name, func(t *testing.T) {
+			other := scope
+			change(&other)
+			setWikiFailureScope(t, other)
+			require.Equal(t, "invalid_configuration", retainWikiIdentityFailure(ctx, scope.ModelID, agent.WikiCandidateSlugPrompt, input, "", raw, original))
+		})
+	}
+	encoded, err := json.Marshal(scope)
+	require.NoError(t, err)
+	for _, setting := range []string{"invalid", string(encoded) + "{}", strings.TrimSuffix(string(encoded), "}") + `,"unknown":true}`, strings.Repeat("x", 4097)} {
+		t.Setenv(wikiFailureCaptureEnv, setting)
+		require.Equal(t, "invalid_configuration", retainWikiIdentityFailure(ctx, scope.ModelID, agent.WikiCandidateSlugPrompt, input, "", raw, original))
+	}
+	files, err := os.ReadDir(scope.Directory)
+	require.NoError(t, err)
+	require.Empty(t, files)
+}
+
+func TestWikiFailureCaptureSourceAndEncodingBounds(t *testing.T) {
+	for _, size := range []int{256 * 1024, 256*1024 + 1} {
+		t.Run(strconv.Itoa(size), func(t *testing.T) {
+			ctx, scope, input, raw, _ := wikiFailureFixture(t)
+			input["Content"] = strings.Repeat("x", size)
+			scope.SourceHash = wikiCheckpointHash([]string{input["Content"]})
+			result, _ := parseWikiCandidateOutput(raw)
+			original := validateWikiFreshSourceIdentities(result, input["Content"])
+			setWikiFailureScope(t, scope)
+			want := "complete"
+			if size > 256*1024 {
+				want = "size_limit"
+			}
+			require.Equal(t, want, retainWikiIdentityFailure(ctx, scope.ModelID, agent.WikiCandidateSlugPrompt, input, "", raw, original))
+			_, artifact := readWikiFailureArtifact(t, scope)
+			require.Equal(t, size, artifact.SourceBytes)
+			if want == "complete" {
+				require.Equal(t, input["Content"], artifact.Source)
+			} else {
+				require.Empty(t, artifact.Source)
+			}
+		})
+	}
+	t.Run("invalid UTF-8 is explicitly unreplayable", func(t *testing.T) {
+		ctx, scope, input, _, original := wikiFailureFixture(t)
+		setWikiFailureScope(t, scope)
+		raw := string([]byte{0xff, 0xfe})
+		require.Equal(t, "invalid_utf8", retainWikiIdentityFailure(ctx, scope.ModelID, agent.WikiCandidateSlugPrompt, input, "", raw, original))
+		_, artifact := readWikiFailureArtifact(t, scope)
+		require.Empty(t, artifact.Response)
+		require.Empty(t, artifact.Source)
+		require.Equal(t, wikiFailureTextHash(raw), artifact.ResponseHash)
+	})
+}
+
+type wikiFailureInjectedFile struct {
+	*os.File
+	failure string
+}
+
+func (f wikiFailureInjectedFile) Write(value []byte) (int, error) {
+	if f.failure == "write" {
+		_, _ = f.File.Write(value[:1])
+		return 1, errors.New("private disk error")
+	}
+	if f.failure == "short" {
+		return 0, nil
+	}
+	return f.File.Write(value)
+}
+func (f wikiFailureInjectedFile) Sync() error {
+	if f.failure == "sync" {
+		return errors.New("private disk error")
+	}
+	return f.File.Sync()
+}
+func (f wikiFailureInjectedFile) Close() error {
+	err := f.File.Close()
+	if f.failure == "close" {
+		return errors.New("private disk error")
+	}
+	return err
+}
+
+func TestWikiFailureCaptureIOFailureConsumesClaim(t *testing.T) {
+	for _, failure := range []string{"open", "write", "short", "sync", "close", "rename"} {
+		t.Run(failure, func(t *testing.T) {
+			_, scope, _, _, original := wikiFailureFixture(t)
+			before := original.Error()
+			outcome := writeWikiFailureArtifactWithFile(scope, []byte(`{"status":"complete"}`), "complete", func(root *os.Root, name string) (wikiFailureFile, error) {
+				if failure == "open" {
+					return nil, errors.New("private disk error")
+				}
+				file, err := root.OpenFile(name, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0600)
+				if failure == "rename" {
+					require.NoError(t, root.Mkdir(filepath.Join(filepath.Dir(name), "failure.json"), 0700))
+				}
+				return wikiFailureInjectedFile{file, failure}, err
+			})
+			require.Equal(t, "write_failed", outcome)
+			require.EqualError(t, original, before)
+			require.Equal(t, "already_claimed", writeWikiFailureArtifact(scope, []byte(`{}`), "complete"))
+			partials, err := filepath.Glob(filepath.Join(scope.Directory, "*", "failure.partial"))
+			require.NoError(t, err)
+			require.Empty(t, partials)
+			completed, err := filepath.Glob(filepath.Join(scope.Directory, "*", "failure.json"))
+			require.NoError(t, err)
+			for _, file := range completed {
+				info, err := os.Stat(file)
+				require.NoError(t, err)
+				require.True(t, info.IsDir())
+			}
+		})
+	}
+}