eric/notarius
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Wire persistent chunk plan caching into the CLI

Browse Source
This commit is contained in:
Eric Rakestraw
2026-07-18 00:31:53 +00:00
parent 51a36efb6b
commit 8ba5228c01
5 changed files with 601 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
docs/roadmap/implementation.md
View File

@@ -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

480
internal/cli/chunk_cache_test.go Normal file
View 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()
}

104
internal/cli/run.go
View File

@@ -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 <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 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" {

21
internal/cli/test_main_test.go Normal file
View 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)
}

21
internal/core/diagnostics/run_dir.go
View File

@@ -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 {