notarius/internal/cli/chunk_cache_test.go

package cli

import (
	"bytes"
	"errors"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"gitea.maximumdirect.net/eric/notarius/internal/core/diagnostics"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
)

type recordingChunkPlanStore struct {
	record   pipeline.ChunkPlanRecord
	decision pipeline.ChunkPlanDecision
	loadErr  error
	saveErr  error
	loads    int
	saves    int
}

func (s *recordingChunkPlanStore) Load(string) (pipeline.ChunkPlanRecord, pipeline.ChunkPlanDecision, error) {
	s.loads++
	return s.record, s.decision, s.loadErr
}

func (s *recordingChunkPlanStore) Save(record pipeline.ChunkPlanRecord) error {
	s.saves++
	s.record = record
	return s.saveErr
}

type recordingChunkPlanFactory struct {
	roots []string
	store *recordingChunkPlanStore
	err   error
}

func (f *recordingChunkPlanFactory) build(root string) (pipeline.ChunkPlanStore, error) {
	f.roots = append(f.roots, root)
	if f.err != nil {
		return nil, f.err
	}
	if f.store == nil {
		f.store = &recordingChunkPlanStore{decision: pipeline.ChunkPlanDecision{Status: pipeline.ChunkPlanMissing}}
	}
	return f.store, nil
}

func TestRunChunkCachePrecedence(t *testing.T) {
	tests := []struct {
		name      string
		fileMode  string
		envMode   string
		flagMode  string
		wantMode  pipeline.ChunkCacheMode
		wantBuild bool
	}{
		{name: "file refresh", fileMode: "refresh", wantMode: pipeline.ChunkCacheRefresh, wantBuild: true},
		{name: "environment over file", fileMode: "bypass", envMode: "refresh", wantMode: pipeline.ChunkCacheRefresh, wantBuild: true},
		{name: "flag over environment", fileMode: "refresh", envMode: "auto", flagMode: "bypass", wantMode: pipeline.ChunkCacheBypass},
		{name: "explicit auto", fileMode: "bypass", envMode: "refresh", flagMode: "auto", wantMode: pipeline.ChunkCacheAuto, wantBuild: true},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			root := filepath.Join(t.TempDir(), "plans")
			configPath := writeTestConfig(t, cacheTestConfig(tc.fileMode, root, ""))
			inputPath := writeFile(t, "input.txt", "input")
			values := map[string]string{}
			if tc.envMode != "" {
				values["NOTARIUS_WORKSPACE_CHUNK_CACHE_MODE"] = tc.envMode
			}
			factory := &recordingChunkPlanFactory{}
			args := []string{"run", "example", "--config", configPath, "--input", inputPath, "--output-dir", t.TempDir()}
			if tc.flagMode != "" {
				args = append(args, "--chunk_cache", tc.flagMode)
			}
			code, stderr := runCacheCommand(t, args, cacheTestOptions(t, values, factory))
			if code != 0 {
				t.Fatalf("code = %d stderr = %q", code, stderr)
			}
			if got := len(factory.roots) > 0; got != tc.wantBuild {
				t.Fatalf("store built = %t roots = %#v, want %t", got, factory.roots, tc.wantBuild)
			}
			if !tc.wantBuild {
				return
			}
			if tc.wantMode == pipeline.ChunkCacheAuto && (factory.store.loads != 1 || factory.store.saves != 1) {
				t.Fatalf("auto calls = load %d save %d", factory.store.loads, factory.store.saves)
			}
			if tc.wantMode == pipeline.ChunkCacheRefresh && (factory.store.loads != 0 || factory.store.saves != 1) {
				t.Fatalf("refresh calls = load %d save %d", factory.store.loads, factory.store.saves)
			}
		})
	}
}

func TestRunChunkCacheInvalidValuesHaveEstablishedExitCodes(t *testing.T) {
	validConfig := writeTestConfig(t, cacheTestConfig("bypass", "", ""))
	invalidFile := writeTestConfig(t, cacheTestConfig("sometimes", "", ""))
	inputPath := writeFile(t, "input.txt", "input")
	tests := []struct {
		name   string
		config string
		env    map[string]string
		flag   string
		want   int
	}{
		{name: "flag", config: validConfig, flag: "sometimes", want: 2},
		{name: "environment", config: validConfig, env: map[string]string{"NOTARIUS_WORKSPACE_CHUNK_CACHE_MODE": "sometimes"}, want: 1},
		{name: "file", config: invalidFile, want: 1},
		{name: "flag does not mask invalid file", config: invalidFile, flag: "bypass", want: 1},
		{name: "flag does not mask invalid environment", config: validConfig, env: map[string]string{"NOTARIUS_WORKSPACE_CHUNK_CACHE_MODE": "sometimes"}, flag: "bypass", want: 1},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			args := []string{"run", "example", "--config", tc.config, "--input", inputPath, "--output-dir", t.TempDir()}
			if tc.flag != "" {
				args = append(args, "--chunk_cache", tc.flag)
			}
			code, _ := runCacheCommand(t, args, cacheTestOptions(t, tc.env, &recordingChunkPlanFactory{}))
			if code != tc.want {
				t.Fatalf("code = %d, want %d", code, tc.want)
			}
		})
	}
}

func TestRunChunkPlanRootResolution(t *testing.T) {
	t.Run("explicit file root", func(t *testing.T) {
		factory := &recordingChunkPlanFactory{}
		resolverCalls := 0
		opts := cacheTestOptions(t, nil, factory)
		opts.UserCacheDir = func() (string, error) { resolverCalls++; return "", errors.New("must not be called") }
		configPath := writeTestConfig(t, cacheTestConfig("auto", "/var/cache/notarius/chunk-plans", ""))
		code, stderr := runCacheCommand(t, cacheRunArgs(t, configPath), opts)
		if code != 0 || stderr != "" || resolverCalls != 0 || len(factory.roots) != 1 || factory.roots[0] != "/var/cache/notarius/chunk-plans" {
			t.Fatalf("code=%d stderr=%q resolver=%d roots=%#v", code, stderr, resolverCalls, factory.roots)
		}
	})

	t.Run("environment root", func(t *testing.T) {
		factory := &recordingChunkPlanFactory{}
		environmentRoot := filepath.Join(t.TempDir(), "environment-plans")
		configPath := writeTestConfig(t, cacheTestConfig("auto", filepath.Join(t.TempDir(), "file-plans"), ""))
		code, stderr := runCacheCommand(t, cacheRunArgs(t, configPath), cacheTestOptions(t, map[string]string{"NOTARIUS_WORKSPACE_CHUNK_CACHE_DIR": environmentRoot}, factory))
		if code != 0 || stderr != "" || len(factory.roots) != 1 || factory.roots[0] != environmentRoot {
			t.Fatalf("code=%d stderr=%q roots=%#v", code, stderr, factory.roots)
		}
	})

	t.Run("per-user default", func(t *testing.T) {
		factory := &recordingChunkPlanFactory{}
		cacheDir := filepath.Join(t.TempDir(), "user-cache")
		opts := cacheTestOptions(t, nil, factory)
		opts.UserCacheDir = func() (string, error) { return cacheDir, nil }
		configPath := writeTestConfig(t, cacheTestConfig("auto", "", ""))
		code, stderr := runCacheCommand(t, cacheRunArgs(t, configPath), opts)
		want := filepath.Join(cacheDir, "notarius", "chunk-plans")
		if code != 0 || stderr != "" || len(factory.roots) != 1 || factory.roots[0] != want {
			t.Fatalf("code=%d stderr=%q roots=%#v want=%q", code, stderr, factory.roots, want)
		}
	})
}

func TestRunBypassSkipsRootAndStore(t *testing.T) {
	factory := &recordingChunkPlanFactory{err: errors.New("must not build")}
	resolverCalls := 0
	opts := cacheTestOptions(t, nil, factory)
	opts.UserCacheDir = func() (string, error) { resolverCalls++; return "", errors.New("must not resolve") }
	configPath := writeTestConfig(t, cacheTestConfig("bypass", "", ""))
	code, stderr := runCacheCommand(t, cacheRunArgs(t, configPath), opts)
	if code != 0 || stderr != "" || resolverCalls != 0 || len(factory.roots) != 0 {
		t.Fatalf("code=%d stderr=%q resolver=%d roots=%#v", code, stderr, resolverCalls, factory.roots)
	}
}

func TestRunChunkPlanSetupFailures(t *testing.T) {
	for _, mode := range []string{"auto", "refresh"} {
		t.Run(mode+" resolver", func(t *testing.T) {
			factory := &recordingChunkPlanFactory{}
			opts := cacheTestOptions(t, nil, factory)
			opts.UserCacheDir = func() (string, error) { return "", errors.New("cache unavailable") }
			configPath := writeTestConfig(t, cacheTestConfig(mode, "", ""))
			code, stderr := runCacheCommand(t, cacheRunArgs(t, configPath), opts)
			if code != 1 || !strings.Contains(stderr, "cache unavailable") || len(factory.roots) != 0 {
				t.Fatalf("code=%d stderr=%q roots=%#v", code, stderr, factory.roots)
			}
		})
	}

	t.Run("store", func(t *testing.T) {
		factory := &recordingChunkPlanFactory{err: errors.New("unwritable store")}
		configPath := writeTestConfig(t, cacheTestConfig("auto", t.TempDir(), ""))
		code, stderr := runCacheCommand(t, cacheRunArgs(t, configPath), cacheTestOptions(t, nil, factory))
		if code != 1 || !strings.Contains(stderr, "unwritable store") {
			t.Fatalf("code=%d stderr=%q", code, stderr)
		}
	})

	t.Run("unwritable filesystem store", func(t *testing.T) {
		root := writeFile(t, "not-a-directory", "occupied")
		configPath := writeTestConfig(t, `version: 2
workspace:
  chunk_cache:
    mode: auto
    directory: `+root+`
pipelines:
  dnd-session:
    input: seriatim
    artifacts:
      spells:
        extract: dnd/spells
`)
		code, stderr := runCacheCommand(t, []string{"run", "dnd-session", "--config", configPath, "--input", writeSeriatimInput(t), "--output-dir", t.TempDir()}, Options{LLMClientFactory: fakeLLMFactory(newFakeRunLLMClient(false), nil), LookupEnv: mapLookup(nil)})
		if code != 1 || (!strings.Contains(stderr, "load chunk plan") && !strings.Contains(stderr, "save chunk plan")) {
			t.Fatalf("code=%d stderr=%q", code, stderr)
		}
	})
}

func TestConfigCommandsDoNotResolveOrCreateChunkPlanState(t *testing.T) {
	configPath := writeTestConfig(t, cacheTestConfig("auto", "", ""))
	for _, args := range [][]string{
		{"config", "validate", "--config", configPath},
		{"pipelines", "list", "--config", configPath},
	} {
		factory := &recordingChunkPlanFactory{err: errors.New("must not build")}
		resolverCalls := 0
		opts := cacheTestOptions(t, nil, factory)
		opts.UserCacheDir = func() (string, error) { resolverCalls++; return "", errors.New("must not resolve") }
		code, stderr := runCacheCommand(t, args, opts)
		if code != 0 || stderr != "" || resolverCalls != 0 || len(factory.roots) != 0 {
			t.Fatalf("args=%v code=%d stderr=%q resolver=%d roots=%#v", args, code, stderr, resolverCalls, factory.roots)
		}
	}
}

func TestRunRecordsExplicitChunkCacheOverrideAndEffectiveMode(t *testing.T) {
	diagnosticsDir := t.TempDir()
	configPath := writeTestConfig(t, cacheTestConfig("bypass", "", diagnosticsDir))
	factory := &recordingChunkPlanFactory{}
	args := append(cacheRunArgs(t, configPath), "--chunk_cache", "refresh")
	code, stderr := runCacheCommand(t, args, cacheTestOptions(t, nil, factory))
	if code != 0 || stderr != "" {
		t.Fatalf("code=%d stderr=%q", code, stderr)
	}
	runDir := onlyChildDir(t, diagnosticsDir)
	invocation := string(readFile(t, filepath.Join(runDir, diagnostics.ArtifactInvocationMetadata)))
	effective := string(readFile(t, filepath.Join(runDir, diagnostics.ArtifactEffectiveConfig)))
	if !strings.Contains(invocation, `"chunk_cache_override": "refresh"`) || !strings.Contains(effective, `"mode": "refresh"`) {
		t.Fatalf("invocation=%s effective=%s", invocation, effective)
	}
}

func TestRunChunkPlanDiagnosticsRedactStoreDecisionReason(t *testing.T) {
	diagnosticsDir := t.TempDir()
	configPath := writeTestConfig(t, cacheTestConfig("auto", t.TempDir(), diagnosticsDir))
	factory := &recordingChunkPlanFactory{store: &recordingChunkPlanStore{
		decision: pipeline.ChunkPlanDecision{Status: pipeline.ChunkPlanInvalid, Reason: "SENTINEL_INVALID_RECORD_CONTENT"},
	}}
	code, stderr := runCacheCommand(t, cacheRunArgs(t, configPath), cacheTestOptions(t, nil, factory))
	if code != 0 || stderr != "" {
		t.Fatalf("code=%d stderr=%q", code, stderr)
	}
	summary := string(readFile(t, filepath.Join(onlyChildDir(t, diagnosticsDir), diagnostics.ArtifactChunkPlan)))
	if strings.Contains(summary, "SENTINEL_INVALID_RECORD_CONTENT") || !strings.Contains(summary, `"lookup_reason": "stored chunk plan is invalid"`) {
		t.Fatalf("chunk plan diagnostic = %s", summary)
	}
}

func TestDefaultAutoReusesPlanAcrossIndependentInvocations(t *testing.T) {
	cacheBase := filepath.Join(t.TempDir(), "cache")
	workspaceDir := filepath.Join(t.TempDir(), "workspace")
	configPath := writeTestConfig(t, `version: 2
workspace:
  directory: `+workspaceDir+`
  debug:
    enabled: true
pipelines:
  dnd-session:
    input: seriatim
    chunk: dnd/scenes
    artifacts:
      spells:
        extract: dnd/spells
`)
	inputPath := writeFile(t, "source.json", `{
  "metadata": {"id": "session-alpha"},
  "segments": [
    {"id": 1, "start": 0, "end": 1, "speaker": "Aria", "text": "Aria casts Cure Wounds."},
    {"id": 2, "start": 1, "end": 2, "speaker": "Borin", "text": "Borin recovers."}
  ]
}`)
	client := newFakeRunLLMClient(false)
	opts := Options{
		LLMClientFactory: fakeLLMFactory(client, nil),
		LookupEnv:        mapLookup(nil),
		UserCacheDir:     func() (string, error) { return cacheBase, nil },
	}
	for i := 0; i < 2; i++ {
		code, stderr := runCacheCommand(t, []string{"run", "dnd-session", "--config", configPath, "--input", inputPath, "--output-dir", t.TempDir()}, opts)
		if code != 0 {
			t.Fatalf("run %d code=%d stderr=%q", i+1, code, stderr)
		}
	}
	if client.calls != 3 {
		t.Fatalf("LLM calls = %d, want chunk+extract then extract-only reuse", client.calls)
	}
	planRoot := filepath.Join(cacheBase, "notarius", "chunk-plans")
	entries, err := os.ReadDir(planRoot)
	if err != nil || len(entries) != 1 {
		t.Fatalf("plan root entries = %v error=%v", entries, err)
	}
	if _, err := os.Stat(filepath.Join(planRoot, entries[0].Name(), "plan.json")); err != nil {
		t.Fatalf("plan file: %v", err)
	}
	debugRuns := childDirs(t, filepath.Join(workspaceDir, "debug"))
	if len(debugRuns) != 2 {
		t.Fatalf("debug runs = %#v", debugRuns)
	}
	attemptCounts := 0
	for _, runDir := range debugRuns {
		if _, err := os.Stat(filepath.Join(runDir, "chunk", "attempt-01.json")); err == nil {
			attemptCounts++
		} else if !os.IsNotExist(err) {
			t.Fatal(err)
		}
	}
	if attemptCounts != 1 {
		t.Fatalf("chunk attempt files across runs = %d, want only generating run", attemptCounts)
	}
}

func TestChunkPlanReuseDependsOnlyOnSourceDigest(t *testing.T) {
	cacheRoot := filepath.Join(t.TempDir(), "plans")
	inputPath := writeSeriatimInput(t)
	referencePath := writeFile(t, "players.txt", "Alyx")
	profilePath := writeScriptoriumProfileFile(t, "chunk-profile", "http://127.0.0.1:8080/v1", "test-model")
	seedConfig := writeTestConfig(t, `version: 2
workspace:
  chunk_cache:
    mode: auto
    directory: `+cacheRoot+`
pipelines:
  seed:
    input: seriatim
    chunk:
      module: generic
      options:
        max_units: 1
        overlap_units: 0
    artifacts:
      spells:
        extract: dnd/spells
`)
	changedConfig := writeTestConfig(t, `version: 2
scriptorium:
  profile_file: `+profilePath+`
workspace:
  chunk_cache:
    mode: auto
    directory: `+cacheRoot+`
pipelines:
  changed:
    input: seriatim
    chunk:
      module: dnd/scenes
      llm_profile: chunk-profile
      references:
        players: `+referencePath+`
      validators:
        - generic/always_accept
    artifacts:
      spells:
        extract: dnd/spells
`)
	client := newFakeRunLLMClient(false)
	opts := Options{LLMClientFactory: fakeLLMFactory(client, nil), LookupEnv: mapLookup(nil)}
	for _, run := range []struct {
		pipeline string
		config   string
	}{{pipeline: "seed", config: seedConfig}, {pipeline: "changed", config: changedConfig}} {
		code, stderr := runCacheCommand(t, []string{"run", run.pipeline, "--config", run.config, "--input", inputPath, "--output-dir", t.TempDir()}, opts)
		if code != 0 {
			t.Fatalf("pipeline %q code=%d stderr=%q", run.pipeline, code, stderr)
		}
	}
	if client.calls != 2 {
		t.Fatalf("LLM calls = %d, want one extractor call per run and no changed chunker call", client.calls)
	}
}

func TestResumeAndChunkCacheModesRemainIndependent(t *testing.T) {
	for _, mode := range []pipeline.ChunkCacheMode{pipeline.ChunkCacheAuto, pipeline.ChunkCacheBypass, pipeline.ChunkCacheRefresh} {
		t.Run(string(mode), func(t *testing.T) {
			workspaceDir := filepath.Join(t.TempDir(), "workspace")
			cacheRoot := filepath.Join(t.TempDir(), "plans")
			configPath := writeTestConfig(t, cacheTestConfigWithWorkspace(string(mode), cacheRoot, workspaceDir))
			factory := &recordingChunkPlanFactory{}
			args := append(cacheRunArgs(t, configPath), "--resume")
			code, stderr := runCacheCommand(t, args, cacheTestOptions(t, nil, factory))
			if code != 0 || stderr != "" {
				t.Fatalf("code=%d stderr=%q", code, stderr)
			}
			if mode == pipeline.ChunkCacheBypass {
				if len(factory.roots) != 0 {
					t.Fatalf("bypass roots = %#v", factory.roots)
				}
			} else if len(factory.roots) != 1 || factory.roots[0] != cacheRoot {
				t.Fatalf("roots = %#v, want %q", factory.roots, cacheRoot)
			}
			if entries := childDirs(t, filepath.Join(workspaceDir, "checkpoints")); len(entries) != 1 {
				t.Fatalf("checkpoint roots = %#v", entries)
			}
			if strings.HasPrefix(cacheRoot, workspaceDir+string(filepath.Separator)) || strings.HasPrefix(workspaceDir, cacheRoot+string(filepath.Separator)) {
				t.Fatalf("cache root %q and workspace root %q overlap", cacheRoot, workspaceDir)
			}
		})
	}
}

func TestWorkspaceDirectoryDoesNotSelectChunkPlanRoot(t *testing.T) {
	cacheBase := filepath.Join(t.TempDir(), "user-cache")
	factory := &recordingChunkPlanFactory{}
	for _, workspaceDir := range []string{filepath.Join(t.TempDir(), "workspace-one"), filepath.Join(t.TempDir(), "workspace-two")} {
		configPath := writeTestConfig(t, cacheTestConfigWithWorkspace("auto", "", workspaceDir))
		opts := cacheTestOptions(t, nil, factory)
		opts.UserCacheDir = func() (string, error) { return cacheBase, nil }
		code, stderr := runCacheCommand(t, cacheRunArgs(t, configPath), opts)
		if code != 0 || stderr != "" {
			t.Fatalf("workspace=%q code=%d stderr=%q", workspaceDir, code, stderr)
		}
	}
	want := filepath.Join(cacheBase, "notarius", "chunk-plans")
	if len(factory.roots) != 2 || factory.roots[0] != want || factory.roots[1] != want {
		t.Fatalf("roots = %#v, want %q twice", factory.roots, want)
	}
}

func cacheTestOptions(t *testing.T, env map[string]string, factory *recordingChunkPlanFactory) Options {
	t.Helper()
	registries := fakeExecutionRegistries(t)
	return Options{
		Catalog:               catalogFromRegistries(registries),
		Registries:            registries,
		LLMClientFactory:      fakeLLMFactory(nil, nil),
		LookupEnv:             mapLookup(env),
		UserCacheDir:          func() (string, error) { return filepath.Join(t.TempDir(), "cache"), nil },
		ChunkPlanStoreFactory: factory.build,
	}
}

func cacheRunArgs(t *testing.T, configPath string) []string {
	t.Helper()
	return []string{"run", "example", "--config", configPath, "--input", writeFile(t, "input.txt", "input"), "--output-dir", t.TempDir()}
}

func runCacheCommand(t *testing.T, args []string, opts Options) (int, string) {
	t.Helper()
	var stdout bytes.Buffer
	var stderr bytes.Buffer
	code := RunWithOptions(args, &stdout, &stderr, opts)
	return code, stderr.String()
}

func cacheTestConfig(mode, directory, diagnosticsDir string) string {
	var b strings.Builder
	b.WriteString("version: 2\n")
	if mode != "" || directory != "" {
		b.WriteString("workspace:\n  chunk_cache:\n")
		if mode != "" {
			b.WriteString("    mode: " + mode + "\n")
		}
		if directory != "" {
			b.WriteString("    directory: " + directory + "\n")
		}
	}
	if diagnosticsDir != "" {
		b.WriteString("diagnostics:\n  work_dir: " + diagnosticsDir + "\n  retention: always\n")
	}
	b.WriteString("pipelines:\n  example:\n    input: fake/input\n    artifacts:\n      spells:\n        extract: fake/extract\n")
	return b.String()
}

func cacheTestConfigWithWorkspace(mode, cacheRoot, workspaceDir string) string {
	var b strings.Builder
	b.WriteString("version: 2\nworkspace:\n  directory: " + workspaceDir + "\n  resume:\n    enabled: true\n  chunk_cache:\n    mode: " + mode + "\n")
	if cacheRoot != "" {
		b.WriteString("    directory: " + cacheRoot + "\n")
	}
	b.WriteString("pipelines:\n  example:\n    input: fake/input\n    artifacts:\n      spells:\n        extract: fake/extract\n")
	return b.String()
}