Wire persistent chunk plan caching into the CLI
This commit is contained in:
@@ -617,7 +617,7 @@ runs are not yet wired to persistent plan storage.
|
|||||||
|
|
||||||
## Stage 5: Wire persistent policy into the CLI
|
## Stage 5: Wire persistent policy into the CLI
|
||||||
|
|
||||||
**Status:** Not started
|
**Status:** Complete
|
||||||
|
|
||||||
### Objective
|
### Objective
|
||||||
|
|
||||||
|
|||||||
480
internal/cli/chunk_cache_test.go
Normal file
480
internal/cli/chunk_cache_test.go
Normal file
@@ -0,0 +1,480 @@
|
|||||||
|
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 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()
|
||||||
|
}
|
||||||
@@ -20,6 +20,7 @@ import (
|
|||||||
"gitea.maximumdirect.net/eric/notarius/internal/core/diagnostics"
|
"gitea.maximumdirect.net/eric/notarius/internal/core/diagnostics"
|
||||||
"gitea.maximumdirect.net/eric/notarius/internal/core/workspace"
|
"gitea.maximumdirect.net/eric/notarius/internal/core/workspace"
|
||||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/checkpoint"
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/checkpoint"
|
||||||
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/chunkplan"
|
||||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
||||||
frameworkdebug "gitea.maximumdirect.net/eric/notarius/internal/framework/debug"
|
frameworkdebug "gitea.maximumdirect.net/eric/notarius/internal/framework/debug"
|
||||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
|
||||||
@@ -30,17 +31,19 @@ const defaultOutputRoot = "./notarius-output"
|
|||||||
|
|
||||||
const usage = `Usage:
|
const usage = `Usage:
|
||||||
notarius help
|
notarius help
|
||||||
notarius run <pipeline-id> --input path/to/source.json [--config path/to/config.yml] [--only lane-a,lane-b] [--resume] [--session-id id] [--reference selector=path] [--without-reference selector]
|
notarius run <pipeline-id> --input path/to/source.json [--config path/to/config.yml] [--only lane-a,lane-b] [--chunk_cache auto|bypass|refresh] [--resume] [--session-id id] [--reference selector=path] [--without-reference selector]
|
||||||
notarius config validate --config path/to/config.yml [--pipeline pipeline-id] [--only lane-a,lane-b]
|
notarius config validate --config path/to/config.yml [--pipeline pipeline-id] [--only lane-a,lane-b]
|
||||||
notarius pipelines list --config path/to/config.yml [--json]
|
notarius pipelines list --config path/to/config.yml [--json]
|
||||||
`
|
`
|
||||||
|
|
||||||
type Options struct {
|
type Options struct {
|
||||||
Catalog pipeline.ModuleCatalog
|
Catalog pipeline.ModuleCatalog
|
||||||
Registries pipeline.Registries
|
Registries pipeline.Registries
|
||||||
LLMClientFactory LLMClientFactory
|
LLMClientFactory LLMClientFactory
|
||||||
LookupEnv func(string) (string, bool)
|
LookupEnv func(string) (string, bool)
|
||||||
Now func() time.Time
|
Now func() time.Time
|
||||||
|
UserCacheDir func() (string, error)
|
||||||
|
ChunkPlanStoreFactory pipeline.ChunkPlanStoreFactory
|
||||||
}
|
}
|
||||||
|
|
||||||
type LLMClientFactory func(ctx context.Context, cfg config.Config, profileID string) (contracts.StructuredLLMClient, []artifacts.LLMProfileManifest, error)
|
type LLMClientFactory func(ctx context.Context, cfg config.Config, profileID string) (contracts.StructuredLLMClient, []artifacts.LLMProfileManifest, error)
|
||||||
@@ -90,6 +93,12 @@ func normalizeOptions(opts Options) (Options, error) {
|
|||||||
if opts.Now == nil {
|
if opts.Now == nil {
|
||||||
opts.Now = time.Now
|
opts.Now = time.Now
|
||||||
}
|
}
|
||||||
|
if opts.UserCacheDir == nil {
|
||||||
|
opts.UserCacheDir = os.UserCacheDir
|
||||||
|
}
|
||||||
|
if opts.ChunkPlanStoreFactory == nil {
|
||||||
|
opts.ChunkPlanStoreFactory = chunkplan.NewFilesystemStore
|
||||||
|
}
|
||||||
if isEmptyCatalog(opts.Catalog) && isEmptyRegistries(opts.Registries) {
|
if isEmptyCatalog(opts.Catalog) && isEmptyRegistries(opts.Registries) {
|
||||||
components, err := newProductionComponents()
|
components, err := newProductionComponents()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -117,10 +126,12 @@ func runPipelineCommand(args []string, stdout, stderr io.Writer, opts Options) i
|
|||||||
diagnosticsDir := fs.String("diagnostics-dir", "", "diagnostics directory")
|
diagnosticsDir := fs.String("diagnostics-dir", "", "diagnostics directory")
|
||||||
llmProfile := fs.String("llm-profile", "", "LLM profile override")
|
llmProfile := fs.String("llm-profile", "", "LLM profile override")
|
||||||
resume := fs.Bool("resume", false, "reuse valid workspace checkpoints")
|
resume := fs.Bool("resume", false, "reuse valid workspace checkpoints")
|
||||||
|
chunkCache := chunkCacheFlag{}
|
||||||
sessionID := sessionIDFlag{}
|
sessionID := sessionIDFlag{}
|
||||||
referenceFlags := stringListFlag{}
|
referenceFlags := stringListFlag{}
|
||||||
withoutReferenceFlags := stringListFlag{}
|
withoutReferenceFlags := stringListFlag{}
|
||||||
fs.Var(&sessionID, "session-id", "prompt session identifier")
|
fs.Var(&sessionID, "session-id", "prompt session identifier")
|
||||||
|
fs.Var(&chunkCache, "chunk_cache", "chunk plan cache mode: auto, bypass, or refresh")
|
||||||
fs.Var(&referenceFlags, "reference", "reference binding, as slot=path, chunk.slot=path, merge.slot=path, lane.slot=path, lane.extract.slot=path, lane.merge.slot=path, or lane.normalize.slot=path")
|
fs.Var(&referenceFlags, "reference", "reference binding, as slot=path, chunk.slot=path, merge.slot=path, lane.slot=path, lane.extract.slot=path, lane.merge.slot=path, or lane.normalize.slot=path")
|
||||||
fs.Var(&withoutReferenceFlags, "without-reference", "unbind a reference, using the same selector forms as --reference")
|
fs.Var(&withoutReferenceFlags, "without-reference", "unbind a reference, using the same selector forms as --reference")
|
||||||
if err := validateRunFlagValues(args); err != nil {
|
if err := validateRunFlagValues(args); err != nil {
|
||||||
@@ -173,6 +184,9 @@ func runPipelineCommand(args []string, stdout, stderr io.Writer, opts Options) i
|
|||||||
fmt.Fprintf(stderr, "notarius: %v\n", err)
|
fmt.Fprintf(stderr, "notarius: %v\n", err)
|
||||||
return 1
|
return 1
|
||||||
}
|
}
|
||||||
|
if chunkCache.set {
|
||||||
|
cfg.Workspace.ChunkCache.Mode = chunkCache.value
|
||||||
|
}
|
||||||
workspaceSettings := workspace.FromConfig(cfg)
|
workspaceSettings := workspace.FromConfig(cfg)
|
||||||
if dir := strings.TrimSpace(*diagnosticsDir); dir != "" {
|
if dir := strings.TrimSpace(*diagnosticsDir); dir != "" {
|
||||||
workspaceSettings.DiagnosticsRoot = dir
|
workspaceSettings.DiagnosticsRoot = dir
|
||||||
@@ -191,15 +205,16 @@ func runPipelineCommand(args []string, stdout, stderr io.Writer, opts Options) i
|
|||||||
runID = runDir.RunID()
|
runID = runDir.RunID()
|
||||||
}
|
}
|
||||||
invocation := diagnostics.InvocationMetadata{
|
invocation := diagnostics.InvocationMetadata{
|
||||||
Operation: "run",
|
Operation: "run",
|
||||||
PipelineID: pipelineID,
|
PipelineID: pipelineID,
|
||||||
InputPath: strings.TrimSpace(*inputPath),
|
InputPath: strings.TrimSpace(*inputPath),
|
||||||
ConfigPath: loadedConfigPath,
|
ConfigPath: loadedConfigPath,
|
||||||
ConfigSource: configSource(*configPath),
|
ConfigSource: configSource(*configPath),
|
||||||
OnlyLanes: append([]string(nil), only...),
|
OnlyLanes: append([]string(nil), only...),
|
||||||
Resume: *resume,
|
ChunkCacheOverride: chunkCache.explicitValue(),
|
||||||
RunID: runID,
|
Resume: *resume,
|
||||||
StartedAt: startedAt,
|
RunID: runID,
|
||||||
|
StartedAt: startedAt,
|
||||||
}
|
}
|
||||||
if err := writeDiagnostics(runDir, func() error { return runDir.WriteInvocationMetadata(invocation) }); err != nil {
|
if err := writeDiagnostics(runDir, func() error { return runDir.WriteInvocationMetadata(invocation) }); err != nil {
|
||||||
return failPipelineCommand(stderr, runDir, cfg.Diagnostics.Retention, fmt.Errorf("write diagnostics invocation metadata: %w", err))
|
return failPipelineCommand(stderr, runDir, cfg.Diagnostics.Retention, fmt.Errorf("write diagnostics invocation metadata: %w", err))
|
||||||
@@ -287,6 +302,10 @@ func runPipelineCommand(args []string, stdout, stderr io.Writer, opts Options) i
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return failPipelineCommand(stderr, runDir, cfg.Diagnostics.Retention, fmt.Errorf("read input %q: %w", strings.TrimSpace(*inputPath), err))
|
return failPipelineCommand(stderr, runDir, cfg.Diagnostics.Retention, fmt.Errorf("read input %q: %w", strings.TrimSpace(*inputPath), err))
|
||||||
}
|
}
|
||||||
|
chunkPlans, err := chunkPlanStoreForRun(effective.Config.Workspace.ChunkCache, opts)
|
||||||
|
if err != nil {
|
||||||
|
return failPipelineCommand(stderr, runDir, cfg.Diagnostics.Retention, err)
|
||||||
|
}
|
||||||
checkpointRecorder, checkpointLoader, err := checkpointHandlersForRun(workspaceSettings, effective.ResolvedPipeline, rawInput, only, llmProfiles, strings.TrimSpace(*llmProfile), strings.TrimSpace(sessionID.value), *resume)
|
checkpointRecorder, checkpointLoader, err := checkpointHandlersForRun(workspaceSettings, effective.ResolvedPipeline, rawInput, only, llmProfiles, strings.TrimSpace(*llmProfile), strings.TrimSpace(sessionID.value), *resume)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return failPipelineCommand(stderr, runDir, cfg.Diagnostics.Retention, err)
|
return failPipelineCommand(stderr, runDir, cfg.Diagnostics.Retention, err)
|
||||||
@@ -302,6 +321,8 @@ func runPipelineCommand(args []string, stdout, stderr io.Writer, opts Options) i
|
|||||||
LLMProfiles: llmProfiles,
|
LLMProfiles: llmProfiles,
|
||||||
Metadata: runMetadata(*outputDir, *diagnosticsDir),
|
Metadata: runMetadata(*outputDir, *diagnosticsDir),
|
||||||
Warnings: referenceWarnings,
|
Warnings: referenceWarnings,
|
||||||
|
ChunkCacheMode: effective.Config.Workspace.ChunkCache.Mode,
|
||||||
|
ChunkPlans: chunkPlans,
|
||||||
Checkpoints: checkpointRecorder,
|
Checkpoints: checkpointRecorder,
|
||||||
Checkpoint: checkpointLoader,
|
Checkpoint: checkpointLoader,
|
||||||
Debug: debugRecorder,
|
Debug: debugRecorder,
|
||||||
@@ -598,13 +619,64 @@ func reorderRunArgs(args []string) []string {
|
|||||||
|
|
||||||
func runFlagTakesValue(arg string) bool {
|
func runFlagTakesValue(arg string) bool {
|
||||||
switch arg {
|
switch arg {
|
||||||
case "--config", "--input", "--only", "--output-dir", "--diagnostics-dir", "--llm-profile", "--session-id", "--reference", "--without-reference":
|
case "--config", "--input", "--only", "--output-dir", "--diagnostics-dir", "--llm-profile", "--session-id", "--chunk_cache", "--reference", "--without-reference":
|
||||||
return true
|
return true
|
||||||
default:
|
default:
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type chunkCacheFlag struct {
|
||||||
|
value pipeline.ChunkCacheMode
|
||||||
|
set bool
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *chunkCacheFlag) String() string {
|
||||||
|
if f == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return string(f.value)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *chunkCacheFlag) Set(raw string) error {
|
||||||
|
mode, err := pipeline.ParseChunkCacheMode(raw)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
f.value = mode
|
||||||
|
f.set = true
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f chunkCacheFlag) explicitValue() string {
|
||||||
|
if !f.set {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return string(f.value)
|
||||||
|
}
|
||||||
|
|
||||||
|
func chunkPlanStoreForRun(cfg config.WorkspaceChunkCacheConfig, opts Options) (pipeline.ChunkPlanStore, error) {
|
||||||
|
if cfg.Mode == pipeline.ChunkCacheBypass {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
root := strings.TrimSpace(cfg.Directory)
|
||||||
|
if root == "" {
|
||||||
|
var err error
|
||||||
|
root, err = workspace.DefaultChunkPlanRoot(opts.UserCacheDir)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("resolve chunk plan root: %w", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
store, err := opts.ChunkPlanStoreFactory(root)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("create chunk plan store at %q: %w", root, err)
|
||||||
|
}
|
||||||
|
if store == nil {
|
||||||
|
return nil, fmt.Errorf("create chunk plan store at %q: factory returned nil", root)
|
||||||
|
}
|
||||||
|
return store, nil
|
||||||
|
}
|
||||||
|
|
||||||
func validateRunFlagValues(args []string) error {
|
func validateRunFlagValues(args []string) error {
|
||||||
for i, arg := range args {
|
for i, arg := range args {
|
||||||
if arg != "--session-id" {
|
if arg != "--session-id" {
|
||||||
|
|||||||
21
internal/cli/test_main_test.go
Normal file
21
internal/cli/test_main_test.go
Normal file
@@ -0,0 +1,21 @@
|
|||||||
|
package cli
|
||||||
|
|
||||||
|
import (
|
||||||
|
"os"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestMain(m *testing.M) {
|
||||||
|
const name = "NOTARIUS_WORKSPACE_CHUNK_CACHE_MODE"
|
||||||
|
previous, existed := os.LookupEnv(name)
|
||||||
|
if err := os.Setenv(name, "bypass"); err != nil {
|
||||||
|
panic(err)
|
||||||
|
}
|
||||||
|
code := m.Run()
|
||||||
|
if existed {
|
||||||
|
_ = os.Setenv(name, previous)
|
||||||
|
} else {
|
||||||
|
_ = os.Unsetenv(name)
|
||||||
|
}
|
||||||
|
os.Exit(code)
|
||||||
|
}
|
||||||
@@ -48,16 +48,17 @@ type RedactedEffectiveConfigPayload interface {
|
|||||||
|
|
||||||
// InvocationMetadata captures non-secret invocation details for diagnostics.
|
// InvocationMetadata captures non-secret invocation details for diagnostics.
|
||||||
type InvocationMetadata struct {
|
type InvocationMetadata struct {
|
||||||
Operation string `json:"operation"`
|
Operation string `json:"operation"`
|
||||||
PipelineID string `json:"pipeline_id,omitempty"`
|
PipelineID string `json:"pipeline_id,omitempty"`
|
||||||
PipelineDigest string `json:"pipeline_digest,omitempty"`
|
PipelineDigest string `json:"pipeline_digest,omitempty"`
|
||||||
Resume bool `json:"resume,omitempty"`
|
Resume bool `json:"resume,omitempty"`
|
||||||
InputPath string `json:"input_path,omitempty"`
|
InputPath string `json:"input_path,omitempty"`
|
||||||
ConfigPath string `json:"config_path,omitempty"`
|
ConfigPath string `json:"config_path,omitempty"`
|
||||||
ConfigSource string `json:"config_source,omitempty"`
|
ConfigSource string `json:"config_source,omitempty"`
|
||||||
OnlyLanes []string `json:"only_lanes,omitempty"`
|
OnlyLanes []string `json:"only_lanes,omitempty"`
|
||||||
RunID string `json:"run_id"`
|
ChunkCacheOverride string `json:"chunk_cache_override,omitempty"`
|
||||||
StartedAt time.Time `json:"started_at"`
|
RunID string `json:"run_id"`
|
||||||
|
StartedAt time.Time `json:"started_at"`
|
||||||
}
|
}
|
||||||
|
|
||||||
func ShouldRetainRunDirectory(input RetentionDecisionInput) bool {
|
func ShouldRetainRunDirectory(input RetentionDecisionInput) bool {
|
||||||
|
|||||||
Reference in New Issue
Block a user