xinmi/internal/application/service/wiki_failure_capture_test.go

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