diff --git a/docs/roadmap/implementation.md b/docs/roadmap/implementation.md index 91b6ce0..d3703e5 100644 --- a/docs/roadmap/implementation.md +++ b/docs/roadmap/implementation.md @@ -617,7 +617,7 @@ runs are not yet wired to persistent plan storage. ## Stage 5: Wire persistent policy into the CLI -**Status:** Not started +**Status:** Complete ### Objective diff --git a/internal/cli/chunk_cache_test.go b/internal/cli/chunk_cache_test.go new file mode 100644 index 0000000..2547d27 --- /dev/null +++ b/internal/cli/chunk_cache_test.go @@ -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() +} diff --git a/internal/cli/run.go b/internal/cli/run.go index 19cab56..9bd8aeb 100644 --- a/internal/cli/run.go +++ b/internal/cli/run.go @@ -20,6 +20,7 @@ import ( "gitea.maximumdirect.net/eric/notarius/internal/core/diagnostics" "gitea.maximumdirect.net/eric/notarius/internal/core/workspace" "gitea.maximumdirect.net/eric/notarius/internal/framework/checkpoint" + "gitea.maximumdirect.net/eric/notarius/internal/framework/chunkplan" "gitea.maximumdirect.net/eric/notarius/internal/framework/contracts" frameworkdebug "gitea.maximumdirect.net/eric/notarius/internal/framework/debug" "gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline" @@ -30,17 +31,19 @@ const defaultOutputRoot = "./notarius-output" const usage = `Usage: notarius help - notarius run --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 --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 pipelines list --config path/to/config.yml [--json] ` type Options struct { - Catalog pipeline.ModuleCatalog - Registries pipeline.Registries - LLMClientFactory LLMClientFactory - LookupEnv func(string) (string, bool) - Now func() time.Time + Catalog pipeline.ModuleCatalog + Registries pipeline.Registries + LLMClientFactory LLMClientFactory + LookupEnv func(string) (string, bool) + 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) @@ -90,6 +93,12 @@ func normalizeOptions(opts Options) (Options, error) { if opts.Now == nil { 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) { components, err := newProductionComponents() 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") llmProfile := fs.String("llm-profile", "", "LLM profile override") resume := fs.Bool("resume", false, "reuse valid workspace checkpoints") + chunkCache := chunkCacheFlag{} sessionID := sessionIDFlag{} referenceFlags := stringListFlag{} withoutReferenceFlags := stringListFlag{} 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(&withoutReferenceFlags, "without-reference", "unbind a reference, using the same selector forms as --reference") 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) return 1 } + if chunkCache.set { + cfg.Workspace.ChunkCache.Mode = chunkCache.value + } workspaceSettings := workspace.FromConfig(cfg) if dir := strings.TrimSpace(*diagnosticsDir); dir != "" { workspaceSettings.DiagnosticsRoot = dir @@ -191,15 +205,16 @@ func runPipelineCommand(args []string, stdout, stderr io.Writer, opts Options) i runID = runDir.RunID() } invocation := diagnostics.InvocationMetadata{ - Operation: "run", - PipelineID: pipelineID, - InputPath: strings.TrimSpace(*inputPath), - ConfigPath: loadedConfigPath, - ConfigSource: configSource(*configPath), - OnlyLanes: append([]string(nil), only...), - Resume: *resume, - RunID: runID, - StartedAt: startedAt, + Operation: "run", + PipelineID: pipelineID, + InputPath: strings.TrimSpace(*inputPath), + ConfigPath: loadedConfigPath, + ConfigSource: configSource(*configPath), + OnlyLanes: append([]string(nil), only...), + ChunkCacheOverride: chunkCache.explicitValue(), + Resume: *resume, + RunID: runID, + StartedAt: startedAt, } 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)) @@ -287,6 +302,10 @@ func runPipelineCommand(args []string, stdout, stderr io.Writer, opts Options) i if err != nil { 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) if err != nil { 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, Metadata: runMetadata(*outputDir, *diagnosticsDir), Warnings: referenceWarnings, + ChunkCacheMode: effective.Config.Workspace.ChunkCache.Mode, + ChunkPlans: chunkPlans, Checkpoints: checkpointRecorder, Checkpoint: checkpointLoader, Debug: debugRecorder, @@ -598,13 +619,64 @@ func reorderRunArgs(args []string) []string { func runFlagTakesValue(arg string) bool { 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 default: 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 { for i, arg := range args { if arg != "--session-id" { diff --git a/internal/cli/test_main_test.go b/internal/cli/test_main_test.go new file mode 100644 index 0000000..7768d03 --- /dev/null +++ b/internal/cli/test_main_test.go @@ -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) +} diff --git a/internal/core/diagnostics/run_dir.go b/internal/core/diagnostics/run_dir.go index e7fcc69..86474eb 100644 --- a/internal/core/diagnostics/run_dir.go +++ b/internal/core/diagnostics/run_dir.go @@ -48,16 +48,17 @@ type RedactedEffectiveConfigPayload interface { // InvocationMetadata captures non-secret invocation details for diagnostics. type InvocationMetadata struct { - Operation string `json:"operation"` - PipelineID string `json:"pipeline_id,omitempty"` - PipelineDigest string `json:"pipeline_digest,omitempty"` - Resume bool `json:"resume,omitempty"` - InputPath string `json:"input_path,omitempty"` - ConfigPath string `json:"config_path,omitempty"` - ConfigSource string `json:"config_source,omitempty"` - OnlyLanes []string `json:"only_lanes,omitempty"` - RunID string `json:"run_id"` - StartedAt time.Time `json:"started_at"` + Operation string `json:"operation"` + PipelineID string `json:"pipeline_id,omitempty"` + PipelineDigest string `json:"pipeline_digest,omitempty"` + Resume bool `json:"resume,omitempty"` + InputPath string `json:"input_path,omitempty"` + ConfigPath string `json:"config_path,omitempty"` + ConfigSource string `json:"config_source,omitempty"` + OnlyLanes []string `json:"only_lanes,omitempty"` + ChunkCacheOverride string `json:"chunk_cache_override,omitempty"` + RunID string `json:"run_id"` + StartedAt time.Time `json:"started_at"` } func ShouldRetainRunDirectory(input RetentionDecisionInput) bool {