diff --git a/internal/framework/pipeline/runner.go b/internal/framework/pipeline/runner.go index 4a82db4..d74123e 100644 --- a/internal/framework/pipeline/runner.go +++ b/internal/framework/pipeline/runner.go @@ -81,6 +81,9 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err if r == nil { return output, fmt.Errorf("runner must not be nil") } + if err := ctx.Err(); err != nil { + return output, err + } if err := validateRunInput(input); err != nil { return output, err } @@ -165,6 +168,10 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err if metadataErr != nil { return failOutput(output), fmt.Errorf("clone input adapter metadata: %w", metadataErr) } + if err := ctx.Err(); err != nil { + _ = checkpoints.SourceFailed(adapter.Key(), err) + return failOutput(output), err + } doc, err = adapter.Parse(ctx, contracts.ParseRequest{ SourceID: input.SourceID, Path: input.Path, @@ -172,6 +179,10 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err LLMProfile: input.pipeline.Input.LLMProfile, Metadata: requestMetadata, }) + if ctxErr := ctx.Err(); ctxErr != nil { + _ = checkpoints.SourceFailed(adapter.Key(), ctxErr) + return failOutput(output), ctxErr + } if err != nil { _ = checkpoints.SourceFailed(adapter.Key(), err) return failOutput(output), fmt.Errorf("parse input with adapter %q: %w", adapter.Key(), err) @@ -281,6 +292,9 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err return failOutput(output), err } } + if err := ctx.Err(); err != nil { + return failOutput(output), err + } if len(output.Rejected) > 0 { output.Manifest.ValidationStatus = "rejected" @@ -294,11 +308,17 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err if err := attachModuleManifestMetadata(&output, "output", encoder); err != nil { return failOutput(output), err } + if err := ctx.Err(); err != nil { + return failOutput(output), err + } outputStarted := time.Now().UTC() evidenceArtifact, evidenceSummary, err := buildOutputEvidenceContext(input.Prepared, doc, output.NormalizeOutputs) if err != nil { return failOutput(output), err } + if err := ctx.Err(); err != nil { + return failOutput(output), err + } if evidenceSummary != nil { if err := writeDebugTimed(debugRecorder, "output/evidence-context.json", debugTimedEnvelope{ Stage: string(StageOutput), @@ -308,6 +328,9 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err }); err != nil { return failOutput(output), fmt.Errorf("write evidence context debug artifact: %w", err) } + if err := ctx.Err(); err != nil { + return failOutput(output), err + } } outputDebugPayload := map[string]any{ "manifest": output.Manifest, @@ -328,10 +351,16 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err }); err != nil { return failOutput(output), fmt.Errorf("write output debug artifact: %w", err) } + if err := ctx.Err(); err != nil { + return failOutput(output), err + } outputMetadata, err := cloneMetadata(input.Metadata) if err != nil { return failOutput(output), fmt.Errorf("clone output encoder metadata: %w", err) } + if err := ctx.Err(); err != nil { + return failOutput(output), err + } encoded, err := encoder.Encode(ctx, contracts.OutputRequest{ Manifest: output.Manifest, NormalizeOutputs: cloneSerializedOutputs(output.NormalizeOutputs), @@ -342,7 +371,9 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err ChunkMap: contracts.CloneSerializedArtifactPointer(acceptedChunkMap), EvidenceContext: contracts.CloneSerializedArtifactPointer(evidenceArtifact), }) - output.Warnings = append(output.Warnings, encoded.Warnings...) + if ctxErr := ctx.Err(); ctxErr != nil { + return failOutput(output), ctxErr + } if err != nil { return failOutput(output), fmt.Errorf("encode output with encoder %q: %w", encoder.Key(), err) } @@ -350,7 +381,6 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err if err != nil { return failOutput(output), fmt.Errorf("validate output files from encoder %q: %w", encoder.Key(), err) } - output.OutputFiles = files if err := writeDebugTimed(debugRecorder, "output/output.json", debugTimedEnvelope{ Stage: string(StageOutput), ModuleKey: encoder.Key(), @@ -362,6 +392,11 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err }); err != nil { return failOutput(output), fmt.Errorf("write output debug artifact: %w", err) } + if err := ctx.Err(); err != nil { + return failOutput(output), err + } + output.Warnings = append(output.Warnings, encoded.Warnings...) + output.OutputFiles = files return output, nil } diff --git a/internal/framework/pipeline/runner_concurrency_test.go b/internal/framework/pipeline/runner_concurrency_test.go index c3cc8a1..02a7418 100644 --- a/internal/framework/pipeline/runner_concurrency_test.go +++ b/internal/framework/pipeline/runner_concurrency_test.go @@ -119,6 +119,96 @@ func (a *countingInputAdapter) Parse(ctx context.Context, request contracts.Pars return a.InputAdapter.Parse(ctx, request) } +type cancelingDebugRecorder struct { + cancel context.CancelFunc + onName string + once sync.Once +} + +func (r *cancelingDebugRecorder) Enabled() bool { return true } + +func (r *cancelingDebugRecorder) WriteJSON(name string, _ any) error { + if name == r.onName { + r.once.Do(r.cancel) + } + return nil +} + +func (*cancelingDebugRecorder) WriteBytes(string, []byte) error { return nil } + +type cancelingOutputEncoder struct { + contracts.OutputEncoder + cancel context.CancelFunc + calls atomic.Int32 +} + +func (e *cancelingOutputEncoder) Encode(context.Context, contracts.OutputRequest) (contracts.OutputResult, error) { + e.calls.Add(1) + e.cancel() + return contracts.OutputResult{ + Files: []contracts.OutputFile{{Name: "result.txt", ContentType: "text/plain", Bytes: []byte("result")}}, + Warnings: []contracts.Warning{{ReasonCode: "returned-after-cancel"}}, + }, nil +} + +func TestRunnerSkipsInputAdapterWhenContextIsAlreadyCanceled(t *testing.T) { + prepared := preparedConcurrentPipeline(t, 1) + adapter := &countingInputAdapter{InputAdapter: prepared.input} + prepared.input = adapter + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := New().Run(ctx, RunInput{Prepared: prepared, RawInput: []byte("input")}) + if !errors.Is(err, context.Canceled) { + t.Fatalf("Run() error = %v, want context.Canceled", err) + } + if got := adapter.calls.Load(); got != 0 { + t.Fatalf("input Parse calls = %d, want none", got) + } +} + +func TestRunnerSkipsOutputEncodingAfterDebugCancellation(t *testing.T) { + prepared := preparedConcurrentPipeline(t, 1) + encoder := &countingOrderedOutput{} + prepared.output = encoder + ctx, cancel := context.WithCancel(context.Background()) + + _, err := New().Run(ctx, RunInput{ + Prepared: prepared, + RawInput: []byte("input"), + Debug: &cancelingDebugRecorder{cancel: cancel, onName: "output/input.json"}, + }) + if !errors.Is(err, context.Canceled) { + t.Fatalf("Run() error = %v, want context.Canceled", err) + } + if got := encoder.calls.Load(); got != 0 { + t.Fatalf("output Encode calls = %d, want none", got) + } +} + +func TestRunnerDiscardsOutputReturnedAfterCancellation(t *testing.T) { + prepared := preparedConcurrentPipeline(t, 1) + ctx, cancel := context.WithCancel(context.Background()) + encoder := &cancelingOutputEncoder{OutputEncoder: prepared.output, cancel: cancel} + prepared.output = encoder + + output, err := New().Run(ctx, RunInput{Prepared: prepared, RawInput: []byte("input")}) + if !errors.Is(err, context.Canceled) { + t.Fatalf("Run() error = %v, want context.Canceled", err) + } + if got := encoder.calls.Load(); got != 1 { + t.Fatalf("output Encode calls = %d, want one", got) + } + if len(output.OutputFiles) != 0 { + t.Fatalf("output files = %#v, want none", output.OutputFiles) + } + for _, warning := range output.Warnings { + if warning.ReasonCode == "returned-after-cancel" { + t.Fatalf("output warnings include encoder warning after cancellation") + } + } +} + func TestRunnerFailsGeneratedHandoffBeforeConsumerExtraction(t *testing.T) { prepared := preparedConcurrentPipeline(t, 1) input := &countingInputAdapter{InputAdapter: prepared.input}