From 8d62973627a3c49819425afd973c8b25fb7c8475 Mon Sep 17 00:00:00 2001 From: Eric Rakestraw Date: Sat, 18 Jul 2026 15:46:19 +0000 Subject: [PATCH] Add run control contract tests --- internal/cli/run_contract_test.go | 444 +++++++++++++++++++++++++++ internal/cli/state_hardening_test.go | 65 +++- 2 files changed, 496 insertions(+), 13 deletions(-) create mode 100644 internal/cli/run_contract_test.go diff --git a/internal/cli/run_contract_test.go b/internal/cli/run_contract_test.go new file mode 100644 index 0000000..adc6cd5 --- /dev/null +++ b/internal/cli/run_contract_test.go @@ -0,0 +1,444 @@ +package cli + +import ( + "bytes" + "context" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + + "gitea.maximumdirect.net/eric/notarius/internal/core/artifacts" + "gitea.maximumdirect.net/eric/notarius/internal/core/config" + "gitea.maximumdirect.net/eric/notarius/internal/framework/contracts" + "gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline" +) + +func TestRunControlsRejectSyntaxWithoutAllocatingState(t *testing.T) { + tests := []struct { + name string + args func(stateTestRoots) []string + }{ + {name: "missing pipeline", args: func(roots stateTestRoots) []string { + return []string{"run", "--config", roots.config, "--input", roots.input} + }}, + {name: "missing input", args: func(roots stateTestRoots) []string { + return []string{"run", "sample", "--config", roots.config} + }}, + {name: "unknown flag", args: func(roots stateTestRoots) []string { + return []string{"run", "sample", "--config", roots.config, "--input", roots.input, "--unknown"} + }}, + {name: "blank output directory", args: func(roots stateTestRoots) []string { + return []string{"run", "sample", "--config", roots.config, "--input", roots.input, "--output-dir", ""} + }}, + {name: "blank debug directory", args: func(roots stateTestRoots) []string { + return []string{"run", "sample", "--config", roots.config, "--input", roots.input, "--debug-dir", ""} + }}, + {name: "debug directory without debug", args: func(roots stateTestRoots) []string { + return []string{"run", "sample", "--config", roots.config, "--input", roots.input, "--debug-dir", filepath.Join(filepath.Dir(roots.debug), "requested-debug")} + }}, + {name: "blank session ID", args: func(roots stateTestRoots) []string { + return []string{"run", "sample", "--config", roots.config, "--input", roots.input, "--session-id", ""} + }}, + {name: "multiple pipeline IDs", args: func(roots stateTestRoots) []string { + return []string{"run", "sample", "extra", "--config", roots.config, "--input", roots.input} + }}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + roots := newStateTestRoots(t) + var stdout, stderr bytes.Buffer + code := RunWithOptions(tt.args(roots), &stdout, &stderr, newStateTestHarness().options()) + if code != 2 || stdout.Len() != 0 || stderr.Len() == 0 { + t.Fatalf("code=%d stdout=%q stderr=%q", code, stdout.String(), stderr.String()) + } + assertAbsent(t, roots.output) + assertAbsent(t, roots.debug) + }) + } +} + +func TestRunValidFailuresClassifyAndReportDebug(t *testing.T) { + tests := []struct { + name string + args func(stateTestRoots) []string + wantError string + wantDebug bool + }{ + {name: "unknown pipeline", args: func(roots stateTestRoots) []string { + return []string{"run", "missing", "--config", roots.config, "--input", roots.input, "--chunk_cache", "bypass"} + }, wantError: `pipeline "missing"`}, + {name: "unknown lane", args: func(roots stateTestRoots) []string { + return []string{"run", "sample", "--config", roots.config, "--input", roots.input, "--only", "missing", "--chunk_cache", "bypass", "--debug"} + }, wantError: `lane "missing"`, wantDebug: true}, + {name: "unreadable input", args: func(roots stateTestRoots) []string { + return []string{"run", "sample", "--config", roots.config, "--input", filepath.Join(filepath.Dir(roots.input), "unreadable.txt"), "--chunk_cache", "bypass", "--debug"} + }, wantError: "read input", wantDebug: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + roots := newStateTestRoots(t) + var stdout, stderr bytes.Buffer + code := RunWithOptions(tt.args(roots), &stdout, &stderr, newStateTestHarness().options()) + if code != 1 || stdout.Len() != 0 || !strings.Contains(stderr.String(), tt.wantError) { + t.Fatalf("code=%d stdout=%q stderr=%q", code, stdout.String(), stderr.String()) + } + if tt.wantDebug { + if !strings.Contains(stderr.String(), "debug=") { + t.Fatalf("stderr=%q, want debug path", stderr.String()) + } + onlyChildDir(t, roots.debug) + } else { + assertAbsent(t, roots.debug) + } + assertAbsent(t, roots.output) + }) + } +} + +func TestRunOnlyExecutesSelectedLanes(t *testing.T) { + roots := newStateTestRoots(t) + data, err := os.ReadFile(roots.config) + if err != nil { + t.Fatal(err) + } + data = bytes.Replace(data, []byte(" output: test/output\n"), []byte(" other:\n extract: test/extract\n output: test/output\n"), 1) + if err := os.WriteFile(roots.config, data, 0o600); err != nil { + t.Fatal(err) + } + harness := newStateTestHarness() + var stdout, stderr bytes.Buffer + code := RunWithOptions([]string{"run", "sample", "--config", roots.config, "--input", roots.input, "--only", "items", "--chunk_cache", "bypass"}, &stdout, &stderr, harness.options()) + if code != 0 || !strings.Contains(stdout.String(), "outputs=1") || stderr.Len() != 0 { + t.Fatalf("code=%d stdout=%q stderr=%q", code, stdout.String(), stderr.String()) + } + harness.mu.Lock() + extractCalls := harness.extractCalls + harness.mu.Unlock() + if extractCalls != 1 { + t.Fatalf("extract calls = %d, want only the selected lane", extractCalls) + } +} + +func TestRunStateRootsHonorEnvironmentFlagsAndDefaults(t *testing.T) { + t.Run("environment roots", func(t *testing.T) { + roots := newStateTestRoots(t) + environmentOutput := filepath.Join(t.TempDir(), "environment-output") + environmentDebug := filepath.Join(t.TempDir(), "environment-debug") + opts := newStateTestHarness().options() + opts.LookupEnv = lookupRunContractEnv(map[string]string{ + "NOTARIUS_OUTPUT_DIR": environmentOutput, + "NOTARIUS_DEBUG_DIR": environmentDebug, + }) + result := runWithStateRoots(t, roots, opts, nil) + if result.code != 0 { + t.Fatalf("code=%d stdout=%q stderr=%q", result.code, result.stdout, result.stderr) + } + assertFile(t, filepath.Join(environmentOutput, filepath.Base(onlyChildDir(t, environmentOutput)), "result.json")) + onlyChildDir(t, environmentDebug) + assertAbsent(t, roots.output) + assertAbsent(t, roots.debug) + }) + + t.Run("command flags override environment", func(t *testing.T) { + roots := newStateTestRoots(t) + environmentOutput := filepath.Join(t.TempDir(), "environment-output") + environmentDebug := filepath.Join(t.TempDir(), "environment-debug") + flagOutput := filepath.Join(t.TempDir(), "flag-output") + flagDebug := filepath.Join(t.TempDir(), "flag-debug") + opts := newStateTestHarness().options() + opts.LookupEnv = lookupRunContractEnv(map[string]string{ + "NOTARIUS_OUTPUT_DIR": environmentOutput, + "NOTARIUS_DEBUG_DIR": environmentDebug, + }) + result := runWithStateRoots(t, roots, opts, []string{"--output-dir", flagOutput, "--debug-dir", flagDebug}) + if result.code != 0 { + t.Fatalf("code=%d stdout=%q stderr=%q", result.code, result.stdout, result.stderr) + } + assertFile(t, filepath.Join(flagOutput, filepath.Base(onlyChildDir(t, flagOutput)), "result.json")) + onlyChildDir(t, flagDebug) + assertAbsent(t, environmentOutput) + assertAbsent(t, environmentDebug) + }) + + t.Run("built-in roots", func(t *testing.T) { + roots := newStateTestRoots(t) + data, err := os.ReadFile(roots.config) + if err != nil { + t.Fatal(err) + } + text := string(data) + text = strings.Replace(text, fmt.Sprintf(" directory: %q\n", roots.output), "", 1) + text = strings.Replace(text, fmt.Sprintf(" directory: %q\n", roots.debug), "", 1) + if err := os.WriteFile(roots.config, []byte(text), 0o600); err != nil { + t.Fatal(err) + } + workDir := t.TempDir() + t.Chdir(workDir) + opts := newStateTestHarness().options() + result := runWithStateRoots(t, roots, opts, nil) + if result.code != 0 { + t.Fatalf("code=%d stdout=%q stderr=%q", result.code, result.stdout, result.stderr) + } + assertFile(t, filepath.Join(workDir, "notarius-output", filepath.Base(onlyChildDir(t, filepath.Join(workDir, "notarius-output"))), "result.json")) + onlyChildDir(t, filepath.Join(workDir, "notarius-debug")) + }) +} + +func TestRunLLMProfileOverrideAndValidationUseInjectedBoundaries(t *testing.T) { + t.Run("one effective profile reaches the factory and modules", func(t *testing.T) { + roots := newStateTestRoots(t) + profileDir := writeRunContractProfiles(t, "override-profile") + prependRunContractConfig(t, roots, fmt.Sprintf("scriptorium:\n profile_dir: %q\n", profileDir)) + harness := newStateTestHarness() + var factoryProfiles []string + opts := harness.options() + opts.LLMClientFactory = func(_ context.Context, _ config.Config, profileID string) (contracts.StructuredLLMClient, []artifacts.LLMProfileManifest, error) { + factoryProfiles = append(factoryProfiles, profileID) + return nil, nil, nil + } + var stdout, stderr bytes.Buffer + code := RunWithOptions([]string{"run", "sample", "--config", roots.config, "--input", roots.input, "--chunk_cache", "bypass", "--llm-profile", "override-profile"}, &stdout, &stderr, opts) + if code != 0 || stderr.Len() != 0 { + t.Fatalf("code=%d stdout=%q stderr=%q", code, stdout.String(), stderr.String()) + } + if len(factoryProfiles) != 1 || factoryProfiles[0] != "override-profile" { + t.Fatalf("factory profiles = %#v, want one override profile", factoryProfiles) + } + harness.mu.Lock() + profiles := append([]string(nil), harness.moduleProfiles...) + harness.mu.Unlock() + if len(profiles) < 4 { + t.Fatalf("module profiles = %#v, want chunk and lane stage requests", profiles) + } + for _, profile := range profiles { + if profile != "override-profile" { + t.Fatalf("module profiles = %#v, want override on every request", profiles) + } + } + }) + + t.Run("validator profile remains distinct", func(t *testing.T) { + roots := newStateTestRoots(t) + profileDir := writeRunContractProfiles(t, "override-profile", "validator-profile") + prependRunContractConfig(t, roots, fmt.Sprintf("scriptorium:\n profile_dir: %q\n", profileDir)) + harness := newStateTestHarness() + var validatorProfiles []string + opts := harness.options() + registerRunContractValidator(t, &opts, &validatorProfiles) + factoryProfiles := []string{} + opts.LLMClientFactory = func(_ context.Context, _ config.Config, profileID string) (contracts.StructuredLLMClient, []artifacts.LLMProfileManifest, error) { + factoryProfiles = append(factoryProfiles, profileID) + return nil, nil, nil + } + var stdout, stderr bytes.Buffer + code := RunWithOptions([]string{"run", "sample", "--config", roots.config, "--input", roots.input, "--chunk_cache", "bypass", "--llm-profile", "override-profile"}, &stdout, &stderr, opts) + if code != 0 || stderr.Len() != 0 { + t.Fatalf("code=%d stdout=%q stderr=%q", code, stdout.String(), stderr.String()) + } + if len(factoryProfiles) != 1 || factoryProfiles[0] != "" { + t.Fatalf("factory profiles = %#v, want one call without a unique profile", factoryProfiles) + } + if len(validatorProfiles) != 1 || validatorProfiles[0] != "validator-profile" { + t.Fatalf("validator profiles = %#v, want configured validator profile", validatorProfiles) + } + }) + + t.Run("unknown profile is rejected without factory access", func(t *testing.T) { + roots := newStateTestRoots(t) + profileDir := writeRunContractProfiles(t, "override-profile") + prependRunContractConfig(t, roots, fmt.Sprintf("scriptorium:\n profile_dir: %q\n", profileDir)) + factoryCalls := 0 + opts := newStateTestHarness().options() + opts.LLMClientFactory = func(context.Context, config.Config, string) (contracts.StructuredLLMClient, []artifacts.LLMProfileManifest, error) { + factoryCalls++ + return nil, nil, nil + } + var stdout, stderr bytes.Buffer + code := RunWithOptions([]string{"run", "sample", "--config", roots.config, "--input", roots.input, "--chunk_cache", "bypass", "--llm-profile", "missing-profile"}, &stdout, &stderr, opts) + if code != 1 || !strings.Contains(stderr.String(), "not configured") || factoryCalls != 0 || stdout.Len() != 0 { + t.Fatalf("code=%d stdout=%q stderr=%q factoryCalls=%d", code, stdout.String(), stderr.String(), factoryCalls) + } + }) +} + +func TestEffectiveLLMProfileIDsAreSortedDeduplicatedAndLLMOnly(t *testing.T) { + resolved := pipeline.ResolvedPipeline{ + Chunk: pipeline.ModuleBinding{LLMProfile: " zeta "}, + ArtifactLanes: []pipeline.ResolvedArtifactLane{ + {Extract: pipeline.ModuleBinding{LLMProfile: "alpha"}, Merge: pipeline.ModuleBinding{LLMProfile: "zeta"}}, + }, + ValidatorChains: []pipeline.ResolvedValidatorChain{{Validators: []pipeline.ResolvedValidator{ + {Binding: pipeline.ModuleBinding{LLMProfile: "deterministic-profile"}, ExecutionClass: contracts.ExecutionClassDeterministic}, + {Binding: pipeline.ModuleBinding{LLMProfile: "beta"}, ExecutionClass: contracts.ExecutionClassLLMBacked}, + }}}, + } + got := effectiveLLMProfileIDs(resolved) + want := []string{"alpha", "beta", "zeta"} + if strings.Join(got, ",") != strings.Join(want, ",") { + t.Fatalf("effective profiles = %#v, want %#v", got, want) + } +} + +func TestRunSessionIDUsesExplicitValueOrSourceDocumentID(t *testing.T) { + for _, tt := range []struct { + name string + args []string + want string + }{ + {name: "source document", want: "source"}, + {name: "explicit trimmed value", args: []string{"--session-id", " explicit-session "}, want: "explicit-session"}, + } { + t.Run(tt.name, func(t *testing.T) { + roots := newStateTestRoots(t) + harness := newStateTestHarness() + args := append([]string{"run", "sample", "--config", roots.config, "--input", roots.input, "--chunk_cache", "bypass"}, tt.args...) + var stdout, stderr bytes.Buffer + code := RunWithOptions(args, &stdout, &stderr, harness.options()) + if code != 0 || stderr.Len() != 0 { + t.Fatalf("code=%d stdout=%q stderr=%q", code, stdout.String(), stderr.String()) + } + harness.mu.Lock() + sessions := append([]string(nil), harness.sessionIDs...) + harness.mu.Unlock() + if len(sessions) < 4 { + t.Fatalf("session IDs = %#v, want all prompt-facing module requests", sessions) + } + for _, session := range sessions { + if session != tt.want { + t.Fatalf("session IDs = %#v, want %q", sessions, tt.want) + } + } + }) + } +} + +func TestRunFactoryAndPreparationFailuresAreProcessFailures(t *testing.T) { + t.Run("LLM factory", func(t *testing.T) { + roots := newStateTestRoots(t) + opts := newStateTestHarness().options() + opts.LLMClientFactory = func(context.Context, config.Config, string) (contracts.StructuredLLMClient, []artifacts.LLMProfileManifest, error) { + return nil, nil, errors.New("injected LLM factory failure") + } + var stdout, stderr bytes.Buffer + code := RunWithOptions([]string{"run", "sample", "--config", roots.config, "--input", roots.input, "--chunk_cache", "bypass"}, &stdout, &stderr, opts) + if code != 1 || !strings.Contains(stderr.String(), "injected LLM factory failure") || stdout.Len() != 0 { + t.Fatalf("code=%d stdout=%q stderr=%q", code, stdout.String(), stderr.String()) + } + }) + + t.Run("pipeline preparation", func(t *testing.T) { + roots := newStateTestRoots(t) + data, err := os.ReadFile(roots.config) + if err != nil { + t.Fatal(err) + } + data = bytes.Replace(data, []byte("extract: test/extract"), []byte("extract: test/failing-extract"), 1) + if err := os.WriteFile(roots.config, data, 0o600); err != nil { + t.Fatal(err) + } + opts := newStateTestHarness().options() + if err := pipeline.RegisterExtractorBuilder(opts.Registries.Extractors, pipeline.ModuleSpec{Key: "test/failing-extract", Stage: pipeline.StageExtract, Requires: []string{"chunks"}, Provides: []string{"artifact"}, ArtifactKind: stateTestArtifactKind}, func(map[string]any) error { return nil }, func(pipeline.BuildRequest) (contracts.Extractor[stateTestArtifact], error) { + return nil, errors.New("injected extractor construction failure") + }); err != nil { + t.Fatal(err) + } + opts.Catalog = catalogFromRegistries(opts.Registries) + var stdout, stderr bytes.Buffer + code := RunWithOptions([]string{"run", "sample", "--config", roots.config, "--input", roots.input, "--chunk_cache", "bypass"}, &stdout, &stderr, opts) + if code != 1 || !strings.Contains(stderr.String(), "injected extractor construction failure") || stdout.Len() != 0 { + t.Fatalf("code=%d stdout=%q stderr=%q", code, stdout.String(), stderr.String()) + } + }) +} + +func TestRunWarningsRemainSuccessfulAndReachDurableSurfaces(t *testing.T) { + roots := newStateTestRoots(t) + harness := newStateTestHarness() + harness.includeWarnings = true + harness.chunkWarnings = []contracts.Warning{{Scope: "chunk", ReasonCode: "contract-warning", Message: "warning retained"}} + var stdout, stderr bytes.Buffer + code := RunWithOptions([]string{"run", "sample", "--config", roots.config, "--input", roots.input, "--chunk_cache", "bypass", "--debug"}, &stdout, &stderr, harness.options()) + if code != 0 || !strings.Contains(stdout.String(), "outputs=1") || !strings.Contains(stderr.String(), "1 warning(s)") { + t.Fatalf("code=%d stdout=%q stderr=%q", code, stdout.String(), stderr.String()) + } + outputPath := filepath.Join(onlyChildDir(t, roots.output), "result.json") + output, err := os.ReadFile(outputPath) + if err != nil || !strings.Contains(string(output), "contract-warning") { + t.Fatalf("durable output = %q, %v", output, err) + } + bundle := onlyChildDir(t, roots.debug) + var warnings []contracts.Warning + readStateTestSummaryJSON(t, bundle, "warnings.json", &warnings) + if len(warnings) != 1 || warnings[0].ReasonCode != "contract-warning" { + t.Fatalf("debug warnings = %#v", warnings) + } +} + +func runWithStateRoots(t *testing.T, roots stateTestRoots, opts Options, extra []string) stateTestResult { + t.Helper() + args := []string{"run", "sample", "--config", roots.config, "--input", roots.input, "--chunk_cache", "bypass", "--debug"} + args = append(args, extra...) + var stdout, stderr bytes.Buffer + return stateTestResult{code: RunWithOptions(args, &stdout, &stderr, opts), stdout: stdout.String(), stderr: stderr.String()} +} + +func lookupRunContractEnv(values map[string]string) func(string) (string, bool) { + return func(name string) (string, bool) { + value, ok := values[name] + return value, ok + } +} + +func prependRunContractConfig(t *testing.T, roots stateTestRoots, prefix string) { + t.Helper() + data, err := os.ReadFile(roots.config) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(roots.config, append([]byte(prefix), data...), 0o600); err != nil { + t.Fatal(err) + } +} + +func writeRunContractProfiles(t *testing.T, ids ...string) string { + t.Helper() + dir := t.TempDir() + for _, id := range ids { + profile := fmt.Sprintf("id: %s\nendpoint: http://127.0.0.1:1/v1\nmodel: %s-model\n", id, id) + if err := os.WriteFile(filepath.Join(dir, id+".yaml"), []byte(profile), 0o600); err != nil { + t.Fatal(err) + } + } + return dir +} + +func registerRunContractValidator(t *testing.T, opts *Options, profiles *[]string) { + t.Helper() + if err := pipeline.RegisterTypedValidatorBuilder(opts.Registries.Validators, stateTestArtifactKind, pipeline.ValidatorSpec{Key: "run-contract-validator", ExecutionClass: contracts.ExecutionClassLLMBacked}, func(map[string]any) error { return nil }, func(pipeline.BuildRequest) (contracts.TypedValidator[stateTestArtifact], error) { + return runContractValidator{profiles: profiles}, nil + }); err != nil { + t.Fatal(err) + } + if err := opts.Registries.ValidatorChains.Register(pipeline.ValidatorChainMapping{Stage: pipeline.StageExtract, Module: "test/extract", Validators: []pipeline.ModuleBinding{{Module: "run-contract-validator", LLMProfile: "validator-profile"}}}); err != nil { + t.Fatal(err) + } + opts.Catalog = catalogFromRegistries(opts.Registries) +} + +type runContractValidator struct { + profiles *[]string +} + +func (v runContractValidator) Name() string { return "run-contract-validator" } + +func (v runContractValidator) ExecutionClass() contracts.ExecutionClass { + return contracts.ExecutionClassLLMBacked +} + +func (v runContractValidator) Validate(_ context.Context, req contracts.TypedValidationRequest[stateTestArtifact]) (contracts.ValidationResult, error) { + *v.profiles = append(*v.profiles, req.LLMProfile) + return contracts.ValidationResult{Approved: true}, nil +} diff --git a/internal/cli/state_hardening_test.go b/internal/cli/state_hardening_test.go index 5da5634..5315be6 100644 --- a/internal/cli/state_hardening_test.go +++ b/internal/cli/state_hardening_test.go @@ -322,6 +322,15 @@ func TestRunUsesOneInjectedIdentityForDebugOutputAndManifest(t *testing.T) { if manifest.RunID != runID { t.Fatalf("manifest run ID = %q, want %q", manifest.RunID, runID) } + wantStartedAt := time.Unix(1, 0).UTC() + if manifest.StartedAt == nil || !manifest.StartedAt.Equal(wantStartedAt) { + t.Fatalf("manifest started at = %v, want %v", manifest.StartedAt, wantStartedAt) + } + var invocation debugbundle.Invocation + readStateTestSummaryJSON(t, debugPath, "invocation.json", &invocation) + if invocation.RunID != runID || !invocation.StartedAt.Equal(wantStartedAt) { + t.Fatalf("debug invocation identity = %#v, want run %q at %v", invocation, runID, wantStartedAt) + } report := readStateTestRunReport(t, debugPath) if !report.Succeeded || report.RunID != runID || report.PipelineID != "sample" || report.OutputPath != outputPath || report.DebugPath != debugPath || report.OutputCount != 1 || report.RejectedCount != 0 || report.WarningCount != 0 || report.ValidationStatus != "approved" { t.Fatalf("success report = %#v", report) @@ -799,11 +808,15 @@ type stateTestHarness struct { runIDCalls uint64 extractErr error chunkWarnings []contracts.Warning + moduleProfiles []string + sessionIDs []string + outputWarnings []contracts.Warning + includeWarnings bool } func newStateTestHarness() *stateTestHarness { return &stateTestHarness{} } func (h *stateTestHarness) options() Options { - registries := pipeline.Registries{Inputs: pipeline.NewInputAdapterRegistry(), Chunkers: pipeline.NewChunkerRegistry(), ArtifactCodecs: pipeline.NewArtifactCodecRegistry(), Extractors: pipeline.NewExtractorRegistry(), Mergers: pipeline.NewMergerRegistry(), Normalizers: pipeline.NewNormalizerRegistry(), Outputs: pipeline.NewOutputEncoderRegistry()} + registries := pipeline.Registries{Inputs: pipeline.NewInputAdapterRegistry(), Chunkers: pipeline.NewChunkerRegistry(), ArtifactCodecs: pipeline.NewArtifactCodecRegistry(), Extractors: pipeline.NewExtractorRegistry(), Mergers: pipeline.NewMergerRegistry(), Normalizers: pipeline.NewNormalizerRegistry(), Validators: pipeline.NewValidatorRegistry(), ValidatorChains: pipeline.NewValidatorChainRegistry(), Outputs: pipeline.NewOutputEncoderRegistry()} if err := pipeline.RegisterArtifactCodec(registries.ArtifactCodecs, stateTestCodec{}); err != nil { panic(err) } @@ -816,13 +829,15 @@ func (h *stateTestHarness) options() Options { if err := pipeline.RegisterExtractor(registries.Extractors, pipeline.ModuleSpec{Key: "test/extract", Stage: pipeline.StageExtract, Requires: []string{"chunks"}, Provides: []string{"artifact"}, ArtifactKind: stateTestArtifactKind}, func() (contracts.Extractor[stateTestArtifact], error) { return stateTestExtractor{h}, nil }); err != nil { panic(err) } - if err := pipeline.RegisterMerger(registries.Mergers, pipeline.ModuleSpec{Key: "test/merge", Stage: pipeline.StageMerge, Requires: []string{"artifact"}, Provides: []string{"merged"}, ArtifactKind: stateTestArtifactKind}, func() (contracts.Merger[stateTestArtifact], error) { return stateTestMerger{}, nil }); err != nil { + if err := pipeline.RegisterMerger(registries.Mergers, pipeline.ModuleSpec{Key: "test/merge", Stage: pipeline.StageMerge, Requires: []string{"artifact"}, Provides: []string{"merged"}, ArtifactKind: stateTestArtifactKind}, func() (contracts.Merger[stateTestArtifact], error) { return stateTestMerger{harness: h}, nil }); err != nil { panic(err) } - if err := pipeline.RegisterNormalizer(registries.Normalizers, pipeline.ModuleSpec{Key: "test/normalize", Stage: pipeline.StageNormalize, Requires: []string{"merged"}, Provides: []string{"normalized"}, ArtifactKind: stateTestArtifactKind}, func() (contracts.Normalizer[stateTestArtifact], error) { return stateTestNormalizer{}, nil }); err != nil { + if err := pipeline.RegisterNormalizer(registries.Normalizers, pipeline.ModuleSpec{Key: "test/normalize", Stage: pipeline.StageNormalize, Requires: []string{"merged"}, Provides: []string{"normalized"}, ArtifactKind: stateTestArtifactKind}, func() (contracts.Normalizer[stateTestArtifact], error) { return stateTestNormalizer{harness: h}, nil }); err != nil { panic(err) } - if err := registries.Outputs.RegisterWithSpec(pipeline.ModuleSpec{Key: "test/output", Stage: pipeline.StageOutput, Requires: []string{"normalized"}, Provides: []string{"output"}}, func() (contracts.OutputEncoder, error) { return stateTestOutput{}, nil }); err != nil { + if err := registries.Outputs.RegisterWithSpec(pipeline.ModuleSpec{Key: "test/output", Stage: pipeline.StageOutput, Requires: []string{"normalized"}, Provides: []string{"output"}}, func() (contracts.OutputEncoder, error) { + return stateTestOutput{harness: h, includeWarnings: h.includeWarnings}, nil + }); err != nil { panic(err) } return Options{Catalog: catalogFromRegistries(registries), Registries: registries, LookupEnv: emptyLookup, Now: func() time.Time { return time.Unix(1, 0) }, RunIDGenerator: func(startedAt time.Time) (string, error) { @@ -847,6 +862,10 @@ type stateTestChunker struct{ harness *stateTestHarness } func (stateTestChunker) Key() string { return "test/chunk" } func (stateTestChunker) ReferenceSlots() []contracts.ReferenceSlot { return nil } func (c stateTestChunker) Plan(_ context.Context, req contracts.ChunkRequest) (contracts.ChunkPlanResult, error) { + c.harness.mu.Lock() + c.harness.moduleProfiles = append(c.harness.moduleProfiles, req.LLMProfile) + c.harness.sessionIDs = append(c.harness.sessionIDs, req.SessionID) + c.harness.mu.Unlock() c.harness.mu.Lock() c.harness.chunkCalls++ c.harness.mu.Unlock() @@ -879,36 +898,56 @@ type stateTestExtractor struct{ harness *stateTestHarness } func (stateTestExtractor) Key() string { return "test/extract" } func (stateTestExtractor) ReferenceSlots() []contracts.ReferenceSlot { return nil } -func (e stateTestExtractor) Extract(context.Context, contracts.TypedExtractionRequest) (contracts.TypedExtractionResult[stateTestArtifact], error) { +func (e stateTestExtractor) Extract(_ context.Context, req contracts.TypedExtractionRequest) (contracts.TypedExtractionResult[stateTestArtifact], error) { e.harness.mu.Lock() defer e.harness.mu.Unlock() e.harness.extractCalls++ + e.harness.moduleProfiles = append(e.harness.moduleProfiles, req.LLMProfile) + e.harness.sessionIDs = append(e.harness.sessionIDs, req.SessionID) if e.harness.extractErr != nil { return contracts.TypedExtractionResult[stateTestArtifact]{}, e.harness.extractErr } return contracts.TypedExtractionResult[stateTestArtifact]{Value: stateTestArtifact{Value: "ok"}}, nil } -type stateTestMerger struct{} +type stateTestMerger struct{ harness *stateTestHarness } func (stateTestMerger) Key() string { return "test/merge" } -func (stateTestMerger) Merge(_ context.Context, req contracts.TypedMergeRequest[stateTestArtifact]) (contracts.TypedMergeResult[stateTestArtifact], error) { +func (m stateTestMerger) Merge(_ context.Context, req contracts.TypedMergeRequest[stateTestArtifact]) (contracts.TypedMergeResult[stateTestArtifact], error) { + m.harness.mu.Lock() + m.harness.moduleProfiles = append(m.harness.moduleProfiles, req.LLMProfile) + m.harness.sessionIDs = append(m.harness.sessionIDs, req.SessionID) + m.harness.mu.Unlock() return contracts.TypedMergeResult[stateTestArtifact]{Value: req.ExtractOutputs[0].Value}, nil } -type stateTestNormalizer struct{} +type stateTestNormalizer struct{ harness *stateTestHarness } func (stateTestNormalizer) Key() string { return "test/normalize" } func (stateTestNormalizer) ReferenceSlots() []contracts.ReferenceSlot { return nil } -func (stateTestNormalizer) Normalize(_ context.Context, req contracts.TypedNormalizeRequest[stateTestArtifact]) (contracts.TypedNormalizeResult[stateTestArtifact], error) { +func (n stateTestNormalizer) Normalize(_ context.Context, req contracts.TypedNormalizeRequest[stateTestArtifact]) (contracts.TypedNormalizeResult[stateTestArtifact], error) { + n.harness.mu.Lock() + n.harness.moduleProfiles = append(n.harness.moduleProfiles, req.LLMProfile) + n.harness.sessionIDs = append(n.harness.sessionIDs, req.SessionID) + n.harness.mu.Unlock() return contracts.TypedNormalizeResult[stateTestArtifact]{Value: req.MergeOutput.Value}, nil } -type stateTestOutput struct{} +type stateTestOutput struct { + harness *stateTestHarness + includeWarnings bool +} -func (stateTestOutput) Key() string { return "test/output" } -func (stateTestOutput) Encode(context.Context, contracts.OutputRequest) (contracts.OutputResult, error) { - return contracts.OutputResult{Files: []contracts.OutputFile{{Name: "result.json", Bytes: []byte("{\"ok\":true}\n")}}}, nil +func (o stateTestOutput) Key() string { return "test/output" } +func (o stateTestOutput) Encode(_ context.Context, req contracts.OutputRequest) (contracts.OutputResult, error) { + o.harness.mu.Lock() + o.harness.outputWarnings = append([]contracts.Warning(nil), req.Warnings...) + o.harness.mu.Unlock() + data := []byte("{\"ok\":true}\n") + if o.includeWarnings && len(req.Warnings) > 0 { + data = []byte(fmt.Sprintf("{\"ok\":true,\"warnings\":%q}\n", req.Warnings[0].ReasonCode)) + } + return contracts.OutputResult{Files: []contracts.OutputFile{{Name: "result.json", Bytes: data}}}, nil } type failingDebugRecorder struct{}