audita/internal/framework/runner/cancellation_test.go

package runner

import (
	"context"
	"errors"
	"net/http"
	"os"
	"path/filepath"
	"sync/atomic"
	"testing"
	"time"

	ck "gitea.maximumdirect.net/eric/ConfigKit"
	"gitea.maximumdirect.net/eric/audita/internal/core/chunking"
	"gitea.maximumdirect.net/eric/audita/internal/core/config"
	"gitea.maximumdirect.net/eric/audita/internal/core/schema"
	"gitea.maximumdirect.net/eric/audita/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/audita/internal/framework/llm"
	"gitea.maximumdirect.net/eric/audita/internal/framework/proposals"
	"gitea.maximumdirect.net/eric/audita/internal/framework/responseschema"
	"gitea.maximumdirect.net/eric/audita/internal/framework/structuredoutput"
	"gitea.maximumdirect.net/eric/audita/internal/framework/validators"
	"gitea.maximumdirect.net/eric/audita/internal/prompts"
	pk "gitea.maximumdirect.net/eric/promptkit"
)

type contextValidator struct {
	validate func(context.Context, validators.Request) (validators.Result, error)
}

func (v contextValidator) Name() string { return "cancel-test" }
func (v contextValidator) Validate(ctx context.Context, req validators.Request) (validators.Result, error) {
	return v.validate(ctx, req)
}

type cancelTransport struct {
	entered chan struct{}
	calls   atomic.Int32
}

func (t *cancelTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	t.calls.Add(1)
	t.entered <- struct{}{}
	<-req.Context().Done()
	return nil, req.Context().Err()
}

func TestValidatorFailureCancelsActiveAndPendingProviderWork(t *testing.T) {
	dir := t.TempDir()
	if err := os.Mkdir(filepath.Join(dir, "profiles"), 0700); err != nil {
		t.Fatal(err)
	}
	main := filepath.Join(dir, "config.yml")
	body := "version: 2\nbackends:\n  shared:\n    endpoint: http://provider.invalid/v1\n    concurrency: 1\n    queue_capacity: 1\nllm:\n  proposal_profile: base\n  transient_retries: 0\n  repair_attempts: 0\n"
	if err := os.WriteFile(main, []byte(body), 0600); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(dir, "profiles/base.yaml"), []byte("profiles:\n  base:\n    backend: shared\n    model: local\n    timeout_seconds: 30\n"), 0600); err != nil {
		t.Fatal(err)
	}
	loaded, err := config.Load(ck.File(main), "")
	if err != nil {
		t.Fatal(err)
	}
	transport := &cancelTransport{entered: make(chan struct{}, 3)}
	runtime, err := llm.NewRuntime(context.Background(), loaded.Config, main, &http.Client{Transport: transport})
	if err != nil {
		t.Fatal(err)
	}
	collector := llm.NewObservationCollector()
	hooks := collector.Hooks()
	started := make(chan struct{}, 2)
	hooks.OnEvent = func(event pk.ExecutionEvent) {
		if event.Kind == pk.ExecutionStarted {
			started <- struct{}{}
		}
	}
	client, err := runtime.NewRoleClient("validation", hooks)
	if err != nil {
		t.Fatal(err)
	}
	origin := errors.New("originating validator failure")
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	validator := contextValidator{validate: func(ctx context.Context, req validators.Request) (validators.Result, error) {
		if *req.SectionIndex == 0 {
			// Hold the failure until one call is active and both peer executions started.
			for i := 0; i < 2; i++ {
				select {
				case <-started:
				case <-ctx.Done():
					return validators.Result{}, ctx.Err()
				}
			}
			select {
			case <-transport.entered:
			case <-ctx.Done():
				return validators.Result{}, ctx.Err()
			}
			return validators.Result{}, origin
		}
		var out map[string]any
		_, err := client.CompleteStructured(ctx, structuredoutput.Request{
			PromptID: prompts.PromptIDValidatorEditorialReview, PromptVersion: prompts.VersionV1,
			Variables: map[string]string{"TranscriptDescriptionBlock": "", "PayloadJSON": "[]"}, Schema: responseschema.ValidatorDecisionSetKey,
			Context: structuredoutput.DiagnosticContext{ModuleKey: req.ModuleKey, ModuleInstance: req.ModuleInstance, SectionIndex: req.SectionIndex},
		}, &out)
		return validators.Result{}, err
	}}
	module := fakeModule{key: "grammar", policy: proposals.ReplacementPolicyRequireUnique, validators: []contracts.Validator{validator}, proposeF: func(req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
		return []proposals.CorrectionProposal{{TargetSegmentID: 1, OriginalText: "hello", CorrectedText: "Hello", Confidence: 1}}, nil
	}}
	sections := []chunking.Section{}
	for i := 0; i < 3; i++ {
		sections = append(sections, chunking.Section{Index: i, Segments: []schema.Segment{{ID: 1, Text: "hello"}}})
	}
	type pipelineOutcome struct {
		result modulePipelineResult
		err    error
	}
	done := make(chan pipelineOutcome, 1)
	go func() {
		result, err := runModulePipeline(ctx, collectSectionProposalsInput{Module: module, Spec: contracts.ModuleRunSpec{ModuleKey: "grammar", InstanceName: "grammar"}, Policy: module.policy, Config: &loaded.Config, Sections: sections})
		done <- pipelineOutcome{result: result, err: err}
	}()
	select {
	case outcome := <-done:
		if !errors.Is(outcome.err, origin) {
			t.Fatalf("lost originating failure: %v", outcome.err)
		}
		if len(outcome.result.WorkUnits) != 3 {
			t.Fatalf("lost work inventory: %+v", outcome.result.WorkUnits)
		}
		if len(outcome.result.Approved) != 0 {
			t.Fatalf("fatal validation returned unchecked candidates: %+v", outcome.result.Approved)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("validator failure did not cancel and join peers")
	}
	if transport.calls.Load() != 1 {
		t.Fatalf("pending call dispatched: %d", transport.calls.Load())
	}
	snapshot := collector.Snapshot()
	if len(snapshot.Calls) != 2 {
		t.Fatalf("partial observations missing: %+v", snapshot)
	}
	for _, call := range snapshot.Calls {
		if call.Finished.IsZero() || call.Failure != structuredoutput.FailureRuntime {
			t.Fatalf("unfinished/cause lost: %+v", call)
		}
	}
	if got := llm.Summarize(snapshot); got.ProviderAttempts != 1 {
		t.Fatalf("attempt count: %+v", got)
	}
}