audita/internal/framework/runner/repair_test.go

package runner

import (
	"context"
	"encoding/json"
	"errors"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"gitea.maximumdirect.net/eric/audita/v3/internal/core/config"
	"gitea.maximumdirect.net/eric/audita/v3/internal/core/schema"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/outcomes"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/proposal_generation"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/proposals"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/structuredoutput"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/validators"
)

type repairClient struct {
	outputs []string
	calls   []structuredoutput.Request
	err     error
	secrets []string
}

func (c *repairClient) CompleteStructured(_ context.Context, req structuredoutput.Request, out any) (structuredoutput.Response, error) {
	if req.Replay != nil {
		copyReplay := *req.Replay
		copyReplay.Turns = append([]structuredoutput.Message(nil), req.Replay.Turns...)
		req.Replay = &copyReplay
	}
	c.calls = append(c.calls, req)
	if c.err != nil {
		return structuredoutput.Response{}, c.err
	}
	value := c.outputs[len(c.calls)-1]
	if err := json.Unmarshal([]byte(value), out); err != nil {
		return structuredoutput.Response{}, err
	}
	return structuredoutput.Response{Prepared: &structuredoutput.PreparedSnapshot{Messages: []structuredoutput.Message{{Role: "user", Content: "synthetic-secret"}}}, Content: json.RawMessage(value), FinalContent: value, RepairInstruction: []structuredoutput.Message{{Role: "user", Content: "instruction"}}}, nil
}

func (c *repairClient) ConfiguredSecrets(context.Context) []string { return c.secrets }

func TestRepairBranchesRetargetAndPreserveOtherApprovals(t *testing.T) {
	client := &repairClient{outputs: []string{`{"proposal_ref":2,"action":"revise","correction":{"id":2,"original_text":"uh ","corrected_text":"","confidence":0.9}}`}}
	candidates := []proposals.EnrichedCorrectionProposal{
		{CorrectionProposal: proposals.CorrectionProposal{TargetSegmentID: 1, OriginalText: "cat", CorrectedText: "dog", Confidence: .9}, ProposalMetadata: proposals.ProposalMetadata{ProposalIndex: 4}},
		{CorrectionProposal: proposals.CorrectionProposal{TargetSegmentID: 0, OriginalText: "x", CorrectedText: "y", Confidence: .9}, ProposalMetadata: proposals.ProposalMetadata{ProposalIndex: 5}},
	}
	input := validateSectionCandidatesInput{Spec: contracts.ModuleRunSpec{ModuleKey: "grammar", RepairClient: client, Settings: config.Stage{ConfidenceThreshold: .8, LLM: config.LLM{CorrectionRepairAttempts: 1}}}, Policy: proposals.ReplacementPolicyRequireUnique, WorkingTranscript: &schema.Transcript{Segments: []schema.Segment{{ID: 1, Text: "cat"}, {ID: 2, Text: "uh hello"}}}, Validators: []contracts.Validator{validators.ProposalShapeValidator{}, validators.NoEffectValidator{}, validators.OriginalTextPresenceValidator{}, validators.ConfidenceThresholdValidator{}, validators.NonEmptyResultingSegmentValidator{}}, SectionEnriched: candidates, Exchange: &structuredoutput.Exchange{PromptID: "proposal", FinalContent: "original response"}}
	initial, err := validateSectionCandidates(context.Background(), input)
	if err != nil {
		t.Fatal(err)
	}
	result, branches, err := repairCandidates(context.Background(), input, []int{1, 9}, initial, "repair")
	if err != nil {
		t.Fatal(err)
	}
	if len(result.approved) != 2 || result.approved[1].ProposalIndex != 5 || result.approved[1].TargetSegmentID != 2 || result.approved[1].CorrectedText != "" || len(branches) != 1 || branches[0].Outcome != "approved" || len(result.filtered) != 0 {
		t.Fatalf("result=%+v branches=%+v", result, branches)
	}
	var payload proposal_generation.RepairRequest
	if err := json.Unmarshal([]byte(client.calls[0].Variables["RepairRequestJSON"]), &payload); err != nil {
		t.Fatal(err)
	}
	if payload.ProposalRef != 2 || payload.Correction.SegmentReference != 9 {
		t.Fatalf("request: %+v", payload)
	}
}

func TestRepairBranchTerminalOutcomes(t *testing.T) {
	for _, tc := range []struct {
		name    string
		outputs []string
		err     error
		budget  int
		outcome string
		calls   int
	}{
		{"withdraw", []string{`{"proposal_ref":1,"action":"withdraw","correction":null}`}, nil, 1, "withdrawn", 1},
		{"exhausted", []string{`{"proposal_ref":1,"action":"revise","correction":{"id":1,"original_text":"cat","corrected_text":"dog","confidence":0.1}}`, `{"proposal_ref":1,"action":"revise","correction":{"id":1,"original_text":"cat","corrected_text":"dog","confidence":0.2}}`}, nil, 2, "exhausted", 2},
		{"no-effect revision", []string{`{"proposal_ref":1,"action":"revise","correction":{"id":1,"original_text":"cat","corrected_text":"cat","confidence":0.9}}`}, nil, 2, "rejected", 1},
		{"wrong reference", []string{`{"proposal_ref":2,"action":"withdraw","correction":null}`}, nil, 1, "unavailable", 1},
		{"provider failure", nil, &structuredoutput.Failure{Class: structuredoutput.FailureRecoverable, Reason: outcomes.ReasonProviderFailed, Message: "provider failed"}, 1, "unavailable", 1},
	} {
		t.Run(tc.name, func(t *testing.T) {
			client := &repairClient{outputs: tc.outputs, err: tc.err}
			candidate := proposals.EnrichedCorrectionProposal{CorrectionProposal: proposals.CorrectionProposal{TargetSegmentID: 1, OriginalText: "cat", CorrectedText: "dog", Confidence: .1}, ProposalMetadata: proposals.ProposalMetadata{ProposalIndex: 0}}
			input := validateSectionCandidatesInput{Spec: contracts.ModuleRunSpec{RepairClient: client, Settings: config.Stage{ConfidenceThreshold: .8, LLM: config.LLM{CorrectionRepairAttempts: tc.budget}}}, Policy: proposals.ReplacementPolicyRequireUnique, WorkingTranscript: &schema.Transcript{Segments: []schema.Segment{{ID: 1, Text: "cat"}}}, Validators: []contracts.Validator{validators.ProposalShapeValidator{}, validators.NoEffectValidator{}, validators.OriginalTextPresenceValidator{}, validators.ConfidenceThresholdValidator{}}, SectionEnriched: []proposals.EnrichedCorrectionProposal{candidate}, Exchange: &structuredoutput.Exchange{PromptID: "proposal"}}
			initial, err := validateSectionCandidates(context.Background(), input)
			if err != nil {
				t.Fatal(err)
			}
			_, branches, err := repairCandidates(context.Background(), input, []int{1}, initial, "repair")
			if err != nil {
				t.Fatal(err)
			}
			if len(branches) != 1 || branches[0].Outcome != tc.outcome || len(client.calls) != tc.calls || branches[0].AttemptCount != tc.calls || len(branches[0].Attempts) != 0 {
				t.Fatalf("branches=%+v calls=%d", branches, len(client.calls))
			}
			if tc.calls == 2 && len(client.calls[1].Replay.Turns) != 2 {
				t.Fatalf("history: %+v", client.calls[1].Replay.Turns)
			}
		})
	}
}

type repairModule struct{ outcomeModule }

func (m repairModule) RepairPromptID() string { return "repair" }

func TestRunnerAppliesRevisionsInOriginalOrder(t *testing.T) {
	candidates := []proposals.CorrectionProposal{{TargetSegmentID: 1, OriginalText: "cat", CorrectedText: "dog", Confidence: .9}, {TargetSegmentID: 1, OriginalText: "cat", CorrectedText: "fox", Confidence: .1}, {TargetSegmentID: 1, OriginalText: "bird", CorrectedText: "owl", Confidence: .1}}
	client := &repairClient{outputs: []string{`{"proposal_ref":2,"action":"revise","correction":{"id":1,"original_text":"cat","corrected_text":"fox","confidence":0.9}}`,
		`{"proposal_ref":3,"action":"revise","correction":{"id":1,"original_text":"bird","corrected_text":"owl","confidence":0.9}}`,
	}}
	module := repairModule{outcomeModule{fakeModule: fakeModule{key: "grammar", policy: proposals.ReplacementPolicyRequireUnique, validators: []contracts.Validator{validators.ProposalShapeValidator{}, validators.NoEffectValidator{}, validators.OriginalTextPresenceValidator{}, validators.ConfidenceThresholdValidator{}}}, proposeResult: func(req contracts.ProposalRequest) (contracts.ProposalResult, error) {
		return contracts.ProposalResult{Proposals: candidates, SegmentReferences: []int{1, 1, 1}, Exchange: &structuredoutput.Exchange{PromptID: "proposal"}}, nil
	}}}
	runner := New(fakeFactory{modules: map[string]contracts.TranscriptModule{"grammar": module}})
	out, err := runner.Run(context.Background(), RunInput{Transcript: &schema.Transcript{Segments: []schema.Segment{{ID: 1, Text: "cat and bird"}}}, ModuleSpecs: []contracts.ModuleRunSpec{{ModuleKey: "grammar", InstanceName: "grammar", RepairClient: client, Settings: config.Stage{ConfidenceThreshold: .8, LLM: config.LLM{CorrectionRepairAttempts: 1}}}}})
	if err != nil {
		t.Fatal(err)
	}
	if out.FinalTranscript.Segments[0].Text != "dog and owl" || len(out.WorkUnits) != 1 || len(out.WorkUnits[0].Repairs) != 2 || out.WorkUnits[0].Repairs[0].Outcome != "approved_skipped" || out.WorkUnits[0].Repairs[1].Outcome != "applied" || out.WorkUnits[0].CorrectionsApplied != 2 || len(out.ModuleResults[0].SkippedChanges) != 1 || out.ModuleResults[0].SkippedChanges[0].ProposalIndex != 1 {
		t.Fatalf("output=%+v", out)
	}
	operations := out.WorkUnits[0].Operations
	if len(operations) != 2 {
		t.Fatalf("repair operations: %+v", operations)
	}
	for _, op := range operations {
		if op.Reference.Kind != outcomes.OperationCorrectionRepair || len(op.RecoveryReasons) != 1 || op.RecoveryReasons[0] != outcomes.RecoveryCorrectionRevised {
			t.Fatalf("repair operation: %+v", op)
		}
	}
}

func TestRepairRevalidatesEarlierApprovalAndRetainsUnrelatedApprovalOnFailure(t *testing.T) {
	candidates := []proposals.EnrichedCorrectionProposal{
		{CorrectionProposal: proposals.CorrectionProposal{TargetSegmentID: 1, OriginalText: "cat", CorrectedText: "dog", Confidence: .9}, ProposalMetadata: proposals.ProposalMetadata{ProposalIndex: 0}},
		{CorrectionProposal: proposals.CorrectionProposal{TargetSegmentID: 1, OriginalText: "bird", CorrectedText: "owl", Confidence: .1}, ProposalMetadata: proposals.ProposalMetadata{ProposalIndex: 1}},
	}
	for _, tc := range []struct {
		name    string
		failure bool
	}{{"revision", false}, {"provider failure", true}} {
		t.Run(tc.name, func(t *testing.T) {
			calls := 0
			check := fakeValidator{name: "custom", validateF: func(req contracts.ValidationRequest) (validators.Result, error) {
				var decisions []validators.Decision
				for _, candidate := range req.CandidateProposal {
					calls++
					decisions = append(decisions, validators.Decision{ProposalIndex: candidate.ProposalIndex, Approved: true, Disposition: validators.DispositionApproved})
				}
				return validators.Result{Decisions: decisions}, nil
			}}
			client := &repairClient{outputs: []string{`{"proposal_ref":2,"action":"revise","correction":{"id":1,"original_text":"bird","corrected_text":"owl","confidence":0.9}}`}}
			if tc.failure {
				client.err = &structuredoutput.Failure{Class: structuredoutput.FailureRecoverable, Reason: outcomes.ReasonProviderFailed, Message: "failed"}
			}
			input := validateSectionCandidatesInput{Spec: contracts.ModuleRunSpec{RepairClient: client, Settings: config.Stage{ConfidenceThreshold: .8, LLM: config.LLM{CorrectionRepairAttempts: 1}}}, Policy: proposals.ReplacementPolicyRequireUnique, WorkingTranscript: &schema.Transcript{Segments: []schema.Segment{{ID: 1, Text: "cat bird"}}}, Validators: []contracts.Validator{check, validators.ConfidenceThresholdValidator{}}, SectionEnriched: candidates, Exchange: &structuredoutput.Exchange{PromptID: "proposal"}}
			initial, err := validateSectionCandidates(context.Background(), input)
			if err != nil {
				t.Fatal(err)
			}
			got, branches, err := repairCandidates(context.Background(), input, []int{1, 1}, initial, "repair")
			if err != nil {
				t.Fatal(err)
			}
			if tc.failure {
				if len(got.approved) != 1 || got.approved[0].ProposalIndex != 0 || branches[0].Outcome != "unavailable" || calls != 2 {
					t.Fatalf("result=%+v calls=%d", got, calls)
				}
			} else if len(got.approved) != 2 || branches[0].Outcome != "approved" || calls != 3 {
				t.Fatalf("result=%+v calls=%d", got, calls)
			}
		})
	}
}

func TestRepairInterruptionRetainsAttempt(t *testing.T) {
	client := &repairClient{err: &structuredoutput.Failure{Class: structuredoutput.FailureCanceled, Reason: outcomes.ReasonCanceled, Message: "canceled", Cause: context.Canceled}}
	candidate := proposals.EnrichedCorrectionProposal{CorrectionProposal: proposals.CorrectionProposal{TargetSegmentID: 1, OriginalText: "cat", CorrectedText: "dog", Confidence: .1}, ProposalMetadata: proposals.ProposalMetadata{ProposalIndex: 0}}
	input := validateSectionCandidatesInput{Spec: contracts.ModuleRunSpec{RepairClient: client, Settings: config.Stage{ConfidenceThreshold: .8, LLM: config.LLM{CorrectionRepairAttempts: 1}}}, Policy: proposals.ReplacementPolicyRequireUnique, WorkingTranscript: &schema.Transcript{Segments: []schema.Segment{{ID: 1, Text: "cat"}}}, Validators: []contracts.Validator{validators.ConfidenceThresholdValidator{}}, SectionEnriched: []proposals.EnrichedCorrectionProposal{candidate}, Exchange: &structuredoutput.Exchange{PromptID: "proposal"}}
	initial, err := validateSectionCandidates(context.Background(), input)
	if err != nil {
		t.Fatal(err)
	}
	_, branches, err := repairCandidates(context.Background(), input, []int{1}, initial, "repair")
	if !errors.Is(err, context.Canceled) || len(branches) != 1 || branches[0].Outcome != "interrupted" || branches[0].AttemptCount != 1 || len(branches[0].Attempts) != 0 {
		t.Fatalf("branches=%+v error=%v", branches, err)
	}
}

func TestRepairArtifactsAreUniqueAndRedacted(t *testing.T) {
	client := &repairClient{secrets: []string{"synthetic-secret"}, outputs: []string{`{"proposal_ref":1,"action":"withdraw","correction":null}`, `{"proposal_ref":2,"action":"withdraw","correction":null}`}}
	candidates := []proposals.EnrichedCorrectionProposal{
		{CorrectionProposal: proposals.CorrectionProposal{TargetSegmentID: 1, OriginalText: "cat", CorrectedText: "dog", Confidence: .1}, ProposalMetadata: proposals.ProposalMetadata{ProposalIndex: 0}},
		{CorrectionProposal: proposals.CorrectionProposal{TargetSegmentID: 1, OriginalText: "cat", CorrectedText: "fox", Confidence: .1}, ProposalMetadata: proposals.ProposalMetadata{ProposalIndex: 1}},
	}
	input := validateSectionCandidatesInput{Spec: contracts.ModuleRunSpec{InstanceName: "grammar_1", RepairClient: client, Settings: config.Stage{ConfidenceThreshold: .8, LLM: config.LLM{CorrectionRepairAttempts: 1}}}, Policy: proposals.ReplacementPolicyRequireUnique, WorkingTranscript: &schema.Transcript{Segments: []schema.Segment{{ID: 1, Text: "cat"}}}, Validators: []contracts.Validator{validators.ConfidenceThresholdValidator{}}, SectionEnriched: candidates, Exchange: &structuredoutput.Exchange{PromptID: "proposal"}, DiagnosticsDir: t.TempDir()}
	initial, err := validateSectionCandidates(context.Background(), input)
	if err != nil {
		t.Fatal(err)
	}
	_, branches, err := repairCandidates(context.Background(), input, []int{1, 1}, initial, "repair")
	if err != nil {
		t.Fatal(err)
	}
	if len(branches) != 2 {
		t.Fatalf("branches: %+v", branches)
	}
	a, b := branches[0].Attempts[0].Artifacts.RequestPayloadPath, branches[1].Attempts[0].Artifacts.RequestPayloadPath
	if a == "" || a == b {
		t.Fatalf("artifact paths: %q %q", a, b)
	}
	for _, path := range []string{a, b} {
		data, err := os.ReadFile(path)
		if err != nil {
			t.Fatal(err)
		}
		if strings.Contains(string(data), "synthetic-secret") || !strings.Contains(string(data), "[REDACTED]") {
			t.Fatalf("unredacted %s: %s", filepath.Base(path), data)
		}
	}
}

func TestRepairArtifactFailureDoesNotChangeOutcome(t *testing.T) {
	client := &repairClient{outputs: []string{`{"proposal_ref":1,"action":"withdraw","correction":null}`}}
	path := filepath.Join(t.TempDir(), "file")
	if err := os.WriteFile(path, []byte("x"), 0600); err != nil {
		t.Fatal(err)
	}
	warnings := 0
	candidate := proposals.EnrichedCorrectionProposal{CorrectionProposal: proposals.CorrectionProposal{TargetSegmentID: 1, OriginalText: "cat", CorrectedText: "dog", Confidence: .1}, ProposalMetadata: proposals.ProposalMetadata{ProposalIndex: 0}}
	input := validateSectionCandidatesInput{Spec: contracts.ModuleRunSpec{InstanceName: "grammar", RepairClient: client, Settings: config.Stage{ConfidenceThreshold: .8, LLM: config.LLM{CorrectionRepairAttempts: 1}}}, Policy: proposals.ReplacementPolicyRequireUnique, WorkingTranscript: &schema.Transcript{Segments: []schema.Segment{{ID: 1, Text: "cat"}}}, Validators: []contracts.Validator{validators.ConfidenceThresholdValidator{}}, SectionEnriched: []proposals.EnrichedCorrectionProposal{candidate}, Exchange: &structuredoutput.Exchange{PromptID: "proposal"}, DiagnosticsDir: path, ReportDiagnosticError: func(error) { warnings++ }}
	initial, err := validateSectionCandidates(context.Background(), input)
	if err != nil {
		t.Fatal(err)
	}
	_, branches, err := repairCandidates(context.Background(), input, []int{1}, initial, "repair")
	if err != nil || len(branches) != 1 || branches[0].Outcome != "withdrawn" || warnings != 1 {
		t.Fatalf("branches=%+v warnings=%d err=%v", branches, warnings, err)
	}
}

func TestNoEffectAndUnavailableOnlyDoNotStartRepair(t *testing.T) {
	client := &repairClient{}
	for _, rejected := range [][]ValidatorRejectedChange{
		{{ProposalIndex: 0, Disposition: validators.DispositionRejected, ReasonCode: validators.ReasonNoEffect}},
		{{ProposalIndex: 0, Disposition: validators.DispositionUnavailable, ReasonCode: "provider_failed"}},
	} {
		candidate := proposals.EnrichedCorrectionProposal{ProposalMetadata: proposals.ProposalMetadata{ProposalIndex: 0}}
		input := validateSectionCandidatesInput{Spec: contracts.ModuleRunSpec{RepairClient: client, Settings: config.Stage{LLM: config.LLM{CorrectionRepairAttempts: 1}}}, SectionEnriched: []proposals.EnrichedCorrectionProposal{candidate}}
		_, branches, err := repairCandidates(context.Background(), input, nil, validateSectionCandidatesResult{rejected: rejected}, "repair")
		if err != nil || len(branches) != 0 || len(client.calls) != 0 {
			t.Fatalf("branches=%+v calls=%d err=%v", branches, len(client.calls), err)
		}
	}
}

func TestRepairFeedbackRedactsValidationCredentials(t *testing.T) {
	client := &repairClient{outputs: []string{`{"proposal_ref":1,"action":"withdraw","correction":null}`}}
	validation := &repairClient{secrets: []string{"validation-secret"}}
	candidate := proposals.EnrichedCorrectionProposal{CorrectionProposal: proposals.CorrectionProposal{TargetSegmentID: 1, OriginalText: "cat", CorrectedText: "dog", Confidence: .9}, ProposalMetadata: proposals.ProposalMetadata{ProposalIndex: 0}}
	input := validateSectionCandidatesInput{Spec: contracts.ModuleRunSpec{InstanceName: "grammar", RepairClient: client, Settings: config.Stage{LLM: config.LLM{CorrectionRepairAttempts: 1}}}, ValidationLLMClient: validation, SectionEnriched: []proposals.EnrichedCorrectionProposal{candidate}, Exchange: &structuredoutput.Exchange{PromptID: "proposal"}, DiagnosticsDir: t.TempDir()}
	initial := validateSectionCandidatesResult{rejected: []ValidatorRejectedChange{{ProposalIndex: 0, ValidatorName: "editorial_review", Disposition: validators.DispositionRejected, ReasonCode: "rejected", Message: "model echoed validation-secret"}}}
	_, branches, err := repairCandidates(context.Background(), input, []int{1}, initial, "repair")
	if err != nil {
		t.Fatal(err)
	}
	if strings.Contains(client.calls[0].Variables["RepairRequestJSON"], "validation-secret") {
		t.Fatalf("credential in feedback: %s", client.calls[0].Variables["RepairRequestJSON"])
	}
	data, err := os.ReadFile(branches[0].Attempts[0].Artifacts.RequestMetadataPath)
	if err != nil {
		t.Fatal(err)
	}
	if strings.Contains(string(data), "validation-secret") {
		t.Fatalf("credential in artifact: %s", data)
	}
}