Make validation cancellation authoritative

engine_test.go

@@ -2583,6 +2583,40 @@ func TestRunReadsSchemaGraphOncePerOperation(t *testing.T) {
 	}
 }
 
+func TestPreparePreservesValidationCancellationIdentity(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	source := &cancelingPublicSchemaFS{
+		FS: fstest.MapFS{
+			"schemas/schema.json": &fstest.MapFile{Data: []byte(`{"type":"object"}`)},
+		},
+		target: "schemas/schema.json",
+		cancel: cancel,
+	}
+	engine, err := promptkit.NewEngine(promptkit.Config{},
+		promptkit.WithPromptFS(publicStructuredPromptFS("schema.cancel.prompt", "schema.json"), "prompts"),
+		promptkit.WithSchemaFS(source, "schemas"),
+		promptkit.WithProfiles(promptkit.Profile{
+			ID: "contract-fast", Endpoint: "http://example.test/v1", Model: "schema-model",
+		}),
+	)
+	if err != nil {
+		t.Fatalf("construct engine: %v", err)
+	}
+
+	prepared, err := engine.Prepare(ctx, promptkit.RunRequest{
+		PromptID: "schema.cancel.prompt",
+		Inputs: map[string]promptkit.ArtifactRef{
+			"transcript": promptkit.Inline("Rin opens the gate."),
+		},
+	})
+	if prepared != nil || !errors.Is(err, promptkit.ErrValidation) || !errors.Is(err, context.Canceled) {
+		t.Fatalf("Prepare() = (%v, %v), want nil result with validation and context identities", prepared, err)
+	}
+	if reads := source.readCount(); reads != 1 {
+		t.Fatalf("schema reads = %d, want 1", reads)
+	}
+}
+
 func TestPreparedStructuredOutputRetainsExactSchemaNumbers(t *testing.T) {
 	const schema = `{
   "type": "number",
@@ -3260,11 +3294,73 @@ type countingSchemaFS struct {
 	reads map[string]int
 }
 
-func (f *countingSchemaFS) ReadFile(name string) ([]byte, error) {
+type cancelingPublicSchemaFS struct {
+	fs.FS
+	mu       sync.Mutex
+	target   string
+	cancel   context.CancelFunc
+	canceled bool
+	reads    int
+}
+
+func (f *cancelingPublicSchemaFS) Open(name string) (fs.File, error) {
+	file, err := f.FS.Open(name)
+	if err != nil || name != f.target {
+		return file, err
+	}
+	return &cancelingPublicSchemaFile{File: file, owner: f}, nil
+}
+
+func (f *cancelingPublicSchemaFS) readCount() int {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return f.reads
+}
+
+type cancelingPublicSchemaFile struct {
+	fs.File
+	owner *cancelingPublicSchemaFS
+}
+
+func (f *cancelingPublicSchemaFile) Read(buffer []byte) (int, error) {
+	n, err := f.File.Read(buffer)
+	f.owner.mu.Lock()
+	f.owner.reads++
+	if !f.owner.canceled {
+		f.owner.canceled = true
+		f.owner.cancel()
+	}
+	f.owner.mu.Unlock()
+	return n, err
+}
+
+func (f *countingSchemaFS) Open(name string) (fs.File, error) {
+	file, err := f.FS.Open(name)
+	if err != nil {
+		return nil, err
+	}
+	return &countingSchemaFile{File: file, name: name, owner: f}, nil
+}
+
+type countingSchemaFile struct {
+	fs.File
+	name    string
+	owner   *countingSchemaFS
+	counted bool
+}
+
+func (f *countingSchemaFile) Read(buffer []byte) (int, error) {
+	if !f.counted {
+		f.counted = true
+		f.owner.recordRead(f.name)
+	}
+	return f.File.Read(buffer)
+}
+
+func (f *countingSchemaFS) recordRead(name string) {
 	f.mu.Lock()
 	f.reads[name]++
 	f.mu.Unlock()
-	return fs.ReadFile(f.FS, name)
 }
 
 func (f *countingSchemaFS) readCount(name string) int {