Make validation cancellation authoritative

internal/validate/cancellation_test.go

@@ -0,0 +1,311 @@
+package validate
+
+import (
+	"context"
+	"errors"
+	"io/fs"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"testing/fstest"
+	"time"
+
+	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
+	"github.com/santhosh-tekuri/jsonschema/v6"
+)
+
+func TestValidationCancellationBeforeWorkDoesNotOpenSchemaSource(t *testing.T) {
+	source := &countingSchemaFS{FS: fstest.MapFS{
+		"schema.json": {Data: []byte(`{"type":"object"}`)},
+	}}
+	validator := NewFSValidator(source, ".").(ValidationPreparer)
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	plan, err := validator.PrepareValidation(ctx, domain.OutputContract{
+		ValidationMode: domain.ValidationJSONSchema,
+		SchemaPath:     "schema.json",
+	})
+	if plan != nil || !errors.Is(err, context.Canceled) {
+		t.Fatalf("PrepareValidation() = (%v, %v), want nil plan and context cancellation", plan, err)
+	}
+	if opens := source.opens.Load(); opens != 0 {
+		t.Fatalf("schema source opens = %d, want 0", opens)
+	}
+}
+
+func TestValidationCancellationBetweenReferencedSchemaReadChunks(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	source := &controlledSchemaFS{
+		FS: fstest.MapFS{
+			"root.json":  {Data: []byte(`{"$ref":"child.json"}`)},
+			"child.json": {Data: []byte(`{"type":"string"}` + strings.Repeat(" ", schemaReadChunkSize*2))},
+		},
+		target: "child.json",
+		cancel: cancel,
+	}
+	validator := NewFSValidator(source, ".").(ValidationPreparer)
+
+	plan, err := validator.PrepareValidation(ctx, domain.OutputContract{
+		ValidationMode: domain.ValidationJSONSchema,
+		SchemaPath:     "root.json",
+	})
+	if plan != nil || !errors.Is(err, context.Canceled) {
+		t.Fatalf("PrepareValidation() = (%v, %v), want nil plan and context cancellation", plan, err)
+	}
+	if reads := source.reads.Load(); reads != 1 {
+		t.Fatalf("controlled child reads = %d, want 1", reads)
+	}
+	if closes := source.closes.Load(); closes != 1 {
+		t.Fatalf("controlled child closes = %d, want 1", closes)
+	}
+}
+
+func TestDecodeCancellationAfterSynchronousCallWins(t *testing.T) {
+	ctx := newCheckpointContext(2)
+	value, err := decodeJSONValue(ctx, []byte(`{"value":1}`))
+	if value != nil || !errors.Is(err, context.Canceled) {
+		t.Fatalf("decodeJSONValue() = (%v, %v), want nil value and context cancellation", value, err)
+	}
+}
+
+func TestCompileCancellationAfterSynchronousCallWins(t *testing.T) {
+	for _, dependencyErr := range []error{nil, errors.New("compile failed")} {
+		ctx, cancel := context.WithCancel(context.Background())
+		compiled := &jsonschema.Schema{}
+		schema, err := compileJSONSchema(ctx, "promptkit-schema:/root.json", func(string) (*jsonschema.Schema, error) {
+			cancel()
+			return compiled, dependencyErr
+		})
+		if schema != nil || !errors.Is(err, context.Canceled) {
+			t.Fatalf("compileJSONSchema() = (%v, %v), want authoritative context cancellation", schema, err)
+		}
+		if dependencyErr != nil && errors.Is(err, dependencyErr) {
+			t.Fatalf("compileJSONSchema() error = %v, dependency error should not win", err)
+		}
+	}
+}
+
+func TestExecutionCancellationAfterSynchronousCallWins(t *testing.T) {
+	for _, dependencyErr := range []error{nil, errors.New("schema mismatch")} {
+		ctx, cancel := context.WithCancel(context.Background())
+		executor := &controlledSchemaExecutor{cancel: cancel, err: dependencyErr}
+		validationErrors, err := executeJSONSchema(ctx, executor, map[string]any{"ok": true})
+		if validationErrors != nil || !errors.Is(err, context.Canceled) {
+			t.Fatalf("executeJSONSchema() = (%v, %v), want no result and context cancellation", validationErrors, err)
+		}
+		if calls := executor.calls.Load(); calls != 1 {
+			t.Fatalf("schema execution calls = %d, want 1", calls)
+		}
+	}
+}
+
+func TestCancellationDoesNotDetachBlockedSchemaRead(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	started := make(chan struct{})
+	release := make(chan struct{})
+	source := &controlledSchemaFS{
+		FS: fstest.MapFS{
+			"schema.json": {Data: []byte(`{"type":"object"}`)},
+		},
+		target:  "schema.json",
+		started: started,
+		release: release,
+	}
+	validator := NewFSValidator(source, ".").(ValidationPreparer)
+	type outcome struct {
+		plan PreparedValidation
+		err  error
+	}
+	result := make(chan outcome, 1)
+	go func() {
+		plan, err := validator.PrepareValidation(ctx, domain.OutputContract{
+			ValidationMode: domain.ValidationJSONSchema,
+			SchemaPath:     "schema.json",
+		})
+		result <- outcome{plan: plan, err: err}
+	}()
+
+	waitForSignal(t, started, "schema read to start")
+	cancel()
+	select {
+	case got := <-result:
+		t.Fatalf("blocked dependency returned before release: (%v, %v)", got.plan, got.err)
+	default:
+	}
+	close(release)
+	got := waitForOutcome(t, result)
+	if got.plan != nil || !errors.Is(got.err, context.Canceled) {
+		t.Fatalf("PrepareValidation() after release = (%v, %v), want nil plan and context cancellation", got.plan, got.err)
+	}
+	if reads := source.reads.Load(); reads != 1 {
+		t.Fatalf("blocked reads = %d, want 1 completed read", reads)
+	}
+}
+
+func TestCancellationDoesNotDetachBlockedSchemaExecution(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	started := make(chan struct{})
+	release := make(chan struct{})
+	executor := &controlledSchemaExecutor{started: started, release: release}
+	result := make(chan error, 1)
+	go func() {
+		_, err := executeJSONSchema(ctx, executor, map[string]any{"ok": true})
+		result <- err
+	}()
+
+	waitForSignal(t, started, "schema execution to start")
+	cancel()
+	select {
+	case err := <-result:
+		t.Fatalf("blocked dependency returned before release: %v", err)
+	default:
+	}
+	close(release)
+	select {
+	case err := <-result:
+		if !errors.Is(err, context.Canceled) {
+			t.Fatalf("executeJSONSchema() after release = %v, want context cancellation", err)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("timed out waiting for released schema execution")
+	}
+	if calls := executor.calls.Load(); calls != 1 {
+		t.Fatalf("schema execution calls = %d, want 1 completed call", calls)
+	}
+}
+
+type countingSchemaFS struct {
+	fs.FS
+	opens atomic.Int32
+}
+
+func (f *countingSchemaFS) Open(name string) (fs.File, error) {
+	f.opens.Add(1)
+	return f.FS.Open(name)
+}
+
+type controlledSchemaFS struct {
+	fs.FS
+	target  string
+	cancel  context.CancelFunc
+	started chan struct{}
+	release chan struct{}
+	once    sync.Once
+	reads   atomic.Int32
+	closes  atomic.Int32
+}
+
+func (f *controlledSchemaFS) Open(name string) (fs.File, error) {
+	file, err := f.FS.Open(name)
+	if err != nil || name != f.target {
+		return file, err
+	}
+	return &controlledSchemaFile{File: file, owner: f}, nil
+}
+
+type controlledSchemaFile struct {
+	fs.File
+	owner *controlledSchemaFS
+}
+
+func (f *controlledSchemaFile) Read(buffer []byte) (int, error) {
+	f.owner.once.Do(func() {
+		if f.owner.started != nil {
+			close(f.owner.started)
+		}
+		if f.owner.release != nil {
+			<-f.owner.release
+		}
+	})
+	n, err := f.File.Read(buffer)
+	f.owner.reads.Add(1)
+	if f.owner.cancel != nil {
+		f.owner.cancel()
+	}
+	return n, err
+}
+
+func (f *controlledSchemaFile) Close() error {
+	f.owner.closes.Add(1)
+	return f.File.Close()
+}
+
+type controlledSchemaExecutor struct {
+	cancel  context.CancelFunc
+	err     error
+	started chan struct{}
+	release chan struct{}
+	calls   atomic.Int32
+}
+
+func (e *controlledSchemaExecutor) Validate(any) error {
+	e.calls.Add(1)
+	if e.started != nil {
+		close(e.started)
+	}
+	if e.release != nil {
+		<-e.release
+	}
+	if e.cancel != nil {
+		e.cancel()
+	}
+	return e.err
+}
+
+type checkpointContext struct {
+	context.Context
+	mu        sync.Mutex
+	remaining int
+	canceled  bool
+	done      chan struct{}
+}
+
+func newCheckpointContext(checksUntilCancel int) *checkpointContext {
+	return &checkpointContext{
+		Context:   context.Background(),
+		remaining: checksUntilCancel,
+		done:      make(chan struct{}),
+	}
+}
+
+func (c *checkpointContext) Done() <-chan struct{} {
+	return c.done
+}
+
+func (c *checkpointContext) Err() error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.canceled {
+		return context.Canceled
+	}
+	c.remaining--
+	if c.remaining == 0 {
+		c.canceled = true
+		close(c.done)
+		return context.Canceled
+	}
+	return nil
+}
+
+func waitForSignal(t *testing.T, signal <-chan struct{}, description string) {
+	t.Helper()
+	select {
+	case <-signal:
+	case <-time.After(2 * time.Second):
+		t.Fatalf("timed out waiting for %s", description)
+	}
+}
+
+func waitForOutcome[T any](t *testing.T, result <-chan T) T {
+	t.Helper()
+	select {
+	case value := <-result:
+		return value
+	case <-time.After(2 * time.Second):
+		t.Fatal("timed out waiting for released dependency")
+		var zero T
+		return zero
+	}
+}

internal/validate/standard_validator.go

@@ -22,6 +22,8 @@ import (
 
 const jsonSchemaDraft2020 = "https://json-schema.org/draft/2020-12/schema"
 
+const schemaReadChunkSize = 64 * 1024
+
 // StandardValidator provides basic, JSON, and JSON Schema output validation.
 type StandardValidator struct {
 	schemaBaseDir string
@@ -51,7 +53,7 @@ func (v *FSValidator) Validate(ctx context.Context, artifact *domain.Artifact, c
 type preparedValidation struct {
 	contract       domain.OutputContract
 	schemaDocument any
-	schema         *jsonschema.Schema
+	schema         schemaExecutor
 }
 
 func (p *preparedValidation) Validate(ctx context.Context, artifact *domain.Artifact) (domain.ValidationResult, error) {
@@ -62,14 +64,11 @@ func (p *preparedValidation) SchemaDocument() any {
 	return p.schemaDocument
 }
 
-func (p *preparedValidation) validateJSONSchema(instance any, _ string) ([]string, error) {
+func (p *preparedValidation) validateJSONSchema(ctx context.Context, instance any, _ string) ([]string, error) {
 	if p.schema == nil {
 		return nil, errors.New("prepared JSON schema is unavailable")
 	}
-	if err := p.schema.Validate(instance); err != nil {
-		return []string{fmt.Sprintf("json schema validation failed: %v", err)}, nil
-	}
-	return nil, nil
+	return executeJSONSchema(ctx, p.schema, instance)
 }
 
 func (v *StandardValidator) PrepareValidation(ctx context.Context, contract domain.OutputContract) (PreparedValidation, error) {
@@ -82,11 +81,11 @@ func (v *StandardValidator) PrepareValidation(ctx context.Context, contract doma
 		return prepared, nil
 	}
 
-	resolvedSchemaPath, err := v.resolveSchemaPath(contract.SchemaPath)
+	resolvedSchemaPath, err := v.resolveSchemaPath(ctx, contract.SchemaPath)
 	if err != nil {
 		return nil, err
 	}
-	schemaDocument, err := loadJSONSchemaFile(resolvedSchemaPath)
+	schemaDocument, err := loadJSONSchemaFile(ctx, resolvedSchemaPath)
 	if err != nil {
 		return nil, fmt.Errorf("failed to compile JSON schema %q: %w", resolvedSchemaPath, err)
 	}
@@ -94,12 +93,24 @@ func (v *StandardValidator) PrepareValidation(ctx context.Context, contract doma
 	if err != nil {
 		return nil, err
 	}
-	compiler := newSchemaCompiler(standardSchemaLoader{root: schemaRoot})
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	compiler := newSchemaCompiler(standardSchemaLoader{ctx: ctx, root: schemaRoot})
 	resourceURL := fileSchemaResourceURL(resolvedSchemaPath)
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
 	if err := compiler.AddResource(resourceURL.String(), schemaDocument); err != nil {
+		if contextErr := ctx.Err(); contextErr != nil {
+			return nil, contextErr
+		}
 		return nil, fmt.Errorf("failed to register JSON schema %q: %w", resolvedSchemaPath, err)
 	}
-	schema, err := compiler.Compile(resourceURL.String())
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	schema, err := compileJSONSchema(ctx, resourceURL.String(), compiler.Compile)
 	if err != nil {
 		return nil, fmt.Errorf("failed to compile JSON schema %q: %w", resolvedSchemaPath, err)
 	}
@@ -122,16 +133,25 @@ func (v *FSValidator) PrepareValidation(ctx context.Context, contract domain.Out
 		return prepared, nil
 	}
 
-	schemaName, schemaDocument, err := v.loadSchemaDocument(contract.SchemaPath)
+	schemaName, schemaDocument, err := v.loadSchemaDocument(ctx, contract.SchemaPath)
 	if err != nil {
 		return nil, err
 	}
 	resourceURL := fsSchemaResourceURL(schemaName)
-	compiler := newSchemaCompiler(fsSchemaLoader{fsys: v.fsys, root: filecatalog.CleanFSRoot(v.root)})
+	compiler := newSchemaCompiler(fsSchemaLoader{ctx: ctx, fsys: v.fsys, root: filecatalog.CleanFSRoot(v.root)})
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
 	if err := compiler.AddResource(resourceURL.String(), schemaDocument); err != nil {
+		if contextErr := ctx.Err(); contextErr != nil {
+			return nil, contextErr
+		}
 		return nil, fmt.Errorf("failed to register JSON schema %q: %w", schemaName, err)
 	}
-	schema, err := compiler.Compile(resourceURL.String())
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	schema, err := compileJSONSchema(ctx, resourceURL.String(), compiler.Compile)
 	if err != nil {
 		return nil, fmt.Errorf("failed to compile JSON schema %q: %w", schemaName, err)
 	}
@@ -144,7 +164,7 @@ func (v *FSValidator) PrepareValidation(ctx context.Context, contract domain.Out
 	return prepared, nil
 }
 
-type schemaValidatorFunc func(instance any, schemaPath string) ([]string, error)
+type schemaValidatorFunc func(ctx context.Context, instance any, schemaPath string) ([]string, error)
 
 func validateArtifact(ctx context.Context, artifact *domain.Artifact, contract domain.OutputContract, validateSchema schemaValidatorFunc) (domain.ValidationResult, error) {
 	select {
@@ -169,7 +189,11 @@ func validateArtifact(ctx context.Context, artifact *domain.Artifact, contract d
 		res.IsValid = true
 		return res, nil
 	case domain.ValidationBasic:
-		if strings.TrimSpace(string(artifact.Body)) == "" {
+		empty := strings.TrimSpace(string(artifact.Body)) == ""
+		if err := ctx.Err(); err != nil {
+			return domain.ValidationResult{}, err
+		}
+		if empty {
 			res.Status = domain.ValidationFailed
 			res.IsValid = false
 			res.Errors = []string{"output is empty"}
@@ -179,7 +203,11 @@ func validateArtifact(ctx context.Context, artifact *domain.Artifact, contract d
 		res.IsValid = true
 		return res, nil
 	case domain.ValidationJSON:
-		if !json.Valid(artifact.Body) {
+		valid := json.Valid(artifact.Body)
+		if err := ctx.Err(); err != nil {
+			return domain.ValidationResult{}, err
+		}
+		if !valid {
 			res.Status = domain.ValidationFailed
 			res.IsValid = false
 			res.Errors = []string{"invalid JSON"}
@@ -189,15 +217,18 @@ func validateArtifact(ctx context.Context, artifact *domain.Artifact, contract d
 		res.IsValid = true
 		return res, nil
 	case domain.ValidationJSONSchema:
-		instance, jsonErr := decodeJSONValue(artifact.Body)
+		instance, jsonErr := decodeJSONValue(ctx, artifact.Body)
 		if jsonErr != nil {
+			if contextErr := ctx.Err(); contextErr != nil {
+				return domain.ValidationResult{}, contextErr
+			}
 			res.Status = domain.ValidationFailed
 			res.IsValid = false
 			res.Errors = []string{fmt.Sprintf("invalid JSON: %v", jsonErr)}
 			return res, nil
 		}
 
-		validationErrors, err := validateSchema(instance, contract.SchemaPath)
+		validationErrors, err := validateSchema(ctx, instance, contract.SchemaPath)
 		if err != nil {
 			return domain.ValidationResult{}, err
 		}
@@ -216,8 +247,8 @@ func validateArtifact(ctx context.Context, artifact *domain.Artifact, contract d
 	}
 }
 
-func (v *StandardValidator) validateJSONSchema(instance any, schemaPath string) ([]string, error) {
-	resolvedSchemaPath, err := v.resolveSchemaPath(schemaPath)
+func (v *StandardValidator) validateJSONSchema(ctx context.Context, instance any, schemaPath string) ([]string, error) {
+	resolvedSchemaPath, err := v.resolveSchemaPath(ctx, schemaPath)
 	if err != nil {
 		return nil, err
 	}
@@ -226,21 +257,21 @@ func (v *StandardValidator) validateJSONSchema(instance any, schemaPath string)
 	if err != nil {
 		return nil, err
 	}
-	compiler := newSchemaCompiler(standardSchemaLoader{root: schemaRoot})
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	compiler := newSchemaCompiler(standardSchemaLoader{ctx: ctx, root: schemaRoot})
 	resourceURL := fileSchemaResourceURL(resolvedSchemaPath)
-	schema, err := compiler.Compile(resourceURL.String())
+	schema, err := compileJSONSchema(ctx, resourceURL.String(), compiler.Compile)
 	if err != nil {
 		return nil, fmt.Errorf("failed to compile JSON schema %q: %w", resolvedSchemaPath, err)
 	}
 
-	if err := schema.Validate(instance); err != nil {
-		return []string{fmt.Sprintf("json schema validation failed: %v", err)}, nil
-	}
-	return nil, nil
+	return executeJSONSchema(ctx, schema, instance)
 }
 
-func (v *FSValidator) validateJSONSchema(instance any, schemaPath string) ([]string, error) {
-	schemaName, schemaDoc, err := v.loadSchemaDocument(schemaPath)
+func (v *FSValidator) validateJSONSchema(ctx context.Context, instance any, schemaPath string) ([]string, error) {
+	schemaName, schemaDoc, err := v.loadSchemaDocument(ctx, schemaPath)
 	if err != nil {
 		return nil, err
 	}
@@ -249,40 +280,60 @@ func (v *FSValidator) validateJSONSchema(instance any, schemaPath string) ([]str
 	if err := validateSchemaDialect(schemaDoc); err != nil {
 		return nil, fmt.Errorf("failed to compile JSON schema %q: %w", schemaName, err)
 	}
-	compiler := newSchemaCompiler(fsSchemaLoader{fsys: v.fsys, root: filecatalog.CleanFSRoot(v.root)})
+	compiler := newSchemaCompiler(fsSchemaLoader{ctx: ctx, fsys: v.fsys, root: filecatalog.CleanFSRoot(v.root)})
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
 	if err := compiler.AddResource(resourceURL.String(), schemaDoc); err != nil {
+		if contextErr := ctx.Err(); contextErr != nil {
+			return nil, contextErr
+		}
 		return nil, fmt.Errorf("failed to register JSON schema %q: %w", schemaName, err)
 	}
-	schema, err := compiler.Compile(resourceURL.String())
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	schema, err := compileJSONSchema(ctx, resourceURL.String(), compiler.Compile)
 	if err != nil {
 		return nil, fmt.Errorf("failed to compile JSON schema %q: %w", schemaName, err)
 	}
 
-	if err := schema.Validate(instance); err != nil {
-		return []string{fmt.Sprintf("json schema validation failed: %v", err)}, nil
-	}
-	return nil, nil
+	return executeJSONSchema(ctx, schema, instance)
 }
 
-func decodeJSONValue(body []byte) (any, error) {
+func decodeJSONValue(ctx context.Context, body []byte) (any, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
 	decoder := json.NewDecoder(bytes.NewReader(body))
 	decoder.UseNumber()
 
 	var value any
-	if err := decoder.Decode(&value); err != nil {
+	decodeErr := decoder.Decode(&value)
+	if err := ctx.Err(); err != nil {
 		return nil, err
 	}
+	if decodeErr != nil {
+		return nil, decodeErr
+	}
 
 	var trailing any
-	if err := decoder.Decode(&trailing); errors.Is(err, io.EOF) {
-		return value, nil
-	} else if err != nil {
+	trailingErr := decoder.Decode(&trailing)
+	if err := ctx.Err(); err != nil {
 		return nil, err
 	}
+	if errors.Is(trailingErr, io.EOF) {
+		return value, nil
+	} else if trailingErr != nil {
+		return nil, trailingErr
+	}
 	return nil, errors.New("multiple JSON values")
 }
 
-func (v *StandardValidator) resolveSchemaPath(schemaPath string) (string, error) {
+func (v *StandardValidator) resolveSchemaPath(ctx context.Context, schemaPath string) (string, error) {
+	if err := ctx.Err(); err != nil {
+		return "", err
+	}
 	if strings.TrimSpace(schemaPath) == "" {
 		return "", errors.New("schema path is required for json_schema validation")
 	}
@@ -291,13 +342,25 @@ func (v *StandardValidator) resolveSchemaPath(schemaPath string) (string, error)
 	if err != nil {
 		return "", err
 	}
+	if err := ctx.Err(); err != nil {
+		return "", err
+	}
 	resolved, err := containedFilesystemPath(root, schemaPath)
 	if err != nil {
 		return "", err
 	}
+	if err := ctx.Err(); err != nil {
+		return "", err
+	}
 	if _, err := os.Stat(resolved); err != nil {
+		if contextErr := ctx.Err(); contextErr != nil {
+			return "", contextErr
+		}
 		return "", fmt.Errorf("failed to access schema file %q: %w", resolved, err)
 	}
+	if err := ctx.Err(); err != nil {
+		return "", err
+	}
 
 	return resolved, nil
 }
@@ -318,28 +381,39 @@ func (v *StandardValidator) schemaRoot() (string, error) {
 	return resolved, nil
 }
 
-func (v *FSValidator) loadSchemaDocument(schemaPath string) (string, any, error) {
-	resolved, err := v.resolveSchemaPath(schemaPath)
+func (v *FSValidator) loadSchemaDocument(ctx context.Context, schemaPath string) (string, any, error) {
+	resolved, err := v.resolveSchemaPath(ctx, schemaPath)
 	if err != nil {
 		return "", nil, err
 	}
 
-	raw, err := fs.ReadFile(v.fsys, resolved)
+	raw, err := readSchemaFile(ctx, func() (fs.File, error) {
+		return v.fsys.Open(resolved)
+	})
 	if err != nil {
 		return "", nil, fmt.Errorf("failed to read schema file %q: %w", resolved, err)
 	}
 
-	doc, err := decodeJSONValue(raw)
+	doc, err := decodeJSONValue(ctx, raw)
 	if err != nil {
 		return "", nil, fmt.Errorf("failed to decode JSON schema %q: %w", resolved, err)
 	}
 	if err := validateSchemaDialect(doc); err != nil {
+		if contextErr := ctx.Err(); contextErr != nil {
+			return "", nil, contextErr
+		}
 		return "", nil, fmt.Errorf("failed to decode JSON schema %q: %w", resolved, err)
 	}
+	if err := ctx.Err(); err != nil {
+		return "", nil, err
+	}
 	return resolved, doc, nil
 }
 
-func (v *FSValidator) resolveSchemaPath(schemaPath string) (string, error) {
+func (v *FSValidator) resolveSchemaPath(ctx context.Context, schemaPath string) (string, error) {
+	if err := ctx.Err(); err != nil {
+		return "", err
+	}
 	if strings.TrimSpace(schemaPath) == "" {
 		return "", errors.New("schema path is required for json_schema validation")
 	}
@@ -350,8 +424,14 @@ func (v *FSValidator) resolveSchemaPath(schemaPath string) (string, error) {
 	cleanRoot := filecatalog.CleanFSRoot(v.root)
 	rootInfo, err := fs.Stat(v.fsys, cleanRoot)
 	if err != nil {
+		if contextErr := ctx.Err(); contextErr != nil {
+			return "", contextErr
+		}
 		return "", fmt.Errorf("failed to access schema source %q: %w", cleanRoot, err)
 	}
+	if err := ctx.Err(); err != nil {
+		return "", err
+	}
 
 	var resolved string
 	if rootInfo.IsDir() {
@@ -372,8 +452,14 @@ func (v *FSValidator) resolveSchemaPath(schemaPath string) (string, error) {
 	}
 
 	if _, err := fs.Stat(v.fsys, resolved); err != nil {
+		if contextErr := ctx.Err(); contextErr != nil {
+			return "", contextErr
+		}
 		return "", fmt.Errorf("failed to access schema file %q: %w", resolved, err)
 	}
+	if err := ctx.Err(); err != nil {
+		return "", err
+	}
 	return resolved, nil
 }
 
@@ -411,6 +497,35 @@ func newSchemaCompiler(loader jsonschema.URLLoader) *jsonschema.Compiler {
 	return compiler
 }
 
+type schemaExecutor interface {
+	Validate(instance any) error
+}
+
+func compileJSONSchema(ctx context.Context, resourceURL string, compile func(string) (*jsonschema.Schema, error)) (*jsonschema.Schema, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	schema, compileErr := compile(resourceURL)
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	return schema, compileErr
+}
+
+func executeJSONSchema(ctx context.Context, schema schemaExecutor, instance any) ([]string, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	validationErr := schema.Validate(instance)
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	if validationErr != nil {
+		return []string{fmt.Sprintf("json schema validation failed: %v", validationErr)}, nil
+	}
+	return nil, nil
+}
+
 func validateSchemaDialect(doc any) error {
 	object, ok := doc.(map[string]any)
 	if !ok {
@@ -431,10 +546,14 @@ func validateSchemaDialect(doc any) error {
 }
 
 type standardSchemaLoader struct {
+	ctx  context.Context
 	root string
 }
 
 func (l standardSchemaLoader) Load(resourceURL string) (any, error) {
+	if err := l.ctx.Err(); err != nil {
+		return nil, err
+	}
 	parsed, err := url.Parse(resourceURL)
 	if err != nil || parsed.Scheme != "file" || parsed.Host != "" || parsed.RawQuery != "" || parsed.Opaque != "" {
 		return nil, fmt.Errorf("schema reference %q is not a contained file reference", resourceURL)
@@ -445,9 +564,15 @@ func (l standardSchemaLoader) Load(resourceURL string) (any, error) {
 	}
 	resolved, err := containedFilesystemPath(l.root, fileName)
 	if err != nil {
+		if contextErr := l.ctx.Err(); contextErr != nil {
+			return nil, contextErr
+		}
 		return nil, err
 	}
-	return loadJSONSchemaFile(resolved)
+	if err := l.ctx.Err(); err != nil {
+		return nil, err
+	}
+	return loadJSONSchemaFile(l.ctx, resolved)
 }
 
 func containedFilesystemPath(root, name string) (string, error) {
@@ -473,27 +598,86 @@ func containedFilesystemPath(root, name string) (string, error) {
 	return candidate, nil
 }
 
-func loadJSONSchemaFile(name string) (any, error) {
-	raw, err := os.ReadFile(name)
+func loadJSONSchemaFile(ctx context.Context, name string) (any, error) {
+	raw, err := readSchemaFile(ctx, func() (fs.File, error) {
+		return os.Open(name)
+	})
 	if err != nil {
 		return nil, err
 	}
-	doc, err := decodeJSONValue(raw)
+	doc, err := decodeJSONValue(ctx, raw)
 	if err != nil {
 		return nil, err
 	}
 	if err := validateSchemaDialect(doc); err != nil {
+		if contextErr := ctx.Err(); contextErr != nil {
+			return nil, contextErr
+		}
+		return nil, err
+	}
+	if err := ctx.Err(); err != nil {
 		return nil, err
 	}
 	return doc, nil
 }
 
+func readSchemaFile(ctx context.Context, open func() (fs.File, error)) ([]byte, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	file, openErr := open()
+	if err := ctx.Err(); err != nil {
+		if file != nil {
+			_ = file.Close()
+		}
+		return nil, err
+	}
+	if openErr != nil {
+		if file != nil {
+			_ = file.Close()
+		}
+		return nil, openErr
+	}
+	if file == nil {
+		return nil, errors.New("schema source returned a nil file")
+	}
+	defer file.Close()
+
+	var contents []byte
+	chunk := make([]byte, schemaReadChunkSize)
+	for {
+		if err := ctx.Err(); err != nil {
+			return nil, err
+		}
+		n, readErr := file.Read(chunk)
+		if n > 0 {
+			contents = append(contents, chunk[:n]...)
+		}
+		if err := ctx.Err(); err != nil {
+			return nil, err
+		}
+		if errors.Is(readErr, io.EOF) {
+			return contents, nil
+		}
+		if readErr != nil {
+			return nil, readErr
+		}
+		if n == 0 {
+			return nil, io.ErrNoProgress
+		}
+	}
+}
+
 type fsSchemaLoader struct {
+	ctx  context.Context
 	fsys fs.FS
 	root string
 }
 
 func (l fsSchemaLoader) Load(resourceURL string) (any, error) {
+	if err := l.ctx.Err(); err != nil {
+		return nil, err
+	}
 	parsed, err := url.Parse(resourceURL)
 	if err != nil {
 		return nil, fmt.Errorf("invalid schema reference %q: %w", resourceURL, err)
@@ -513,21 +697,35 @@ func (l fsSchemaLoader) Load(resourceURL string) (any, error) {
 
 	rootInfo, err := fs.Stat(l.fsys, l.root)
 	if err != nil {
+		if contextErr := l.ctx.Err(); contextErr != nil {
+			return nil, contextErr
+		}
+		return nil, err
+	}
+	if err := l.ctx.Err(); err != nil {
 		return nil, err
 	}
 	if !rootInfo.IsDir() && name != l.root {
 		return nil, fmt.Errorf("schema reference %q is outside the configured schema file", resourceURL)
 	}
 
-	raw, err := fs.ReadFile(l.fsys, name)
+	raw, err := readSchemaFile(l.ctx, func() (fs.File, error) {
+		return l.fsys.Open(name)
+	})
 	if err != nil {
 		return nil, err
 	}
-	doc, err := decodeJSONValue(raw)
+	doc, err := decodeJSONValue(l.ctx, raw)
 	if err != nil {
 		return nil, err
 	}
 	if err := validateSchemaDialect(doc); err != nil {
+		if contextErr := l.ctx.Err(); contextErr != nil {
+			return nil, contextErr
+		}
+		return nil, err
+	}
+	if err := l.ctx.Err(); err != nil {
 		return nil, err
 	}
 	return doc, nil

internal/validate/standard_validator_test.go

@@ -924,7 +924,7 @@ func TestJSONSchemaDialectIsDraft2020(t *testing.T) {
 func assertSchemaDocument(t *testing.T, got any, expectedJSON []byte) {
 	t.Helper()
 
-	expected, err := decodeJSONValue(expectedJSON)
+	expected, err := decodeJSONValue(context.Background(), expectedJSON)
 	if err != nil {
 		t.Fatalf("decode expected schema document: %v", err)
 	}

internal/validate/validator.go

@@ -7,11 +7,16 @@ import (
 )
 
 // Validator validates the generated artifact based on the output contract.
+// Validation checks ctx around Promptkit-controlled work and synchronous
+// dependency calls. A dependency call already in progress cannot be preempted;
+// after it returns, cancellation takes precedence over its result.
 type Validator interface {
 	Validate(ctx context.Context, artifact *domain.Artifact, contract domain.OutputContract) (domain.ValidationResult, error)
 }
 
 // PreparedValidation validates artifacts against one frozen output contract.
+// Its cancellation boundary is synchronous: Validate does not detach schema
+// execution, and an observed context error prevents publication of a result.
 type PreparedValidation interface {
 	Validate(ctx context.Context, artifact *domain.Artifact) (domain.ValidationResult, error)
 	// SchemaDocument returns the root JSON Schema document used for provider
@@ -21,6 +26,9 @@ type PreparedValidation interface {
 }
 
 // ValidationPreparer freezes validation resources for one output contract.
+// Preparation reads schemas in context-checked chunks and checks ctx around
+// decoding and compilation. Filesystem and compiler calls remain synchronous,
+// so cancellation becomes authoritative when an in-progress call returns.
 type ValidationPreparer interface {
 	PrepareValidation(ctx context.Context, contract domain.OutputContract) (PreparedValidation, error)
 }