notarius/internal/framework/contracts/correction_test.go

package contracts

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strings"
	"testing"
)

func TestSemanticCorrectionOwnsValidatedContent(t *testing.T) {
	assistant := []byte(`{"items":["original"]}`)
	correction, err := NewSemanticCorrection(assistant, "Return one corrected replacement.")
	if err != nil {
		t.Fatalf("NewSemanticCorrection() error = %v", err)
	}
	assistant[0] = '['
	if got := string(correction.AssistantResponse); got != `{"items":["original"]}` {
		t.Fatalf("assistant response = %q, want owned original content", got)
	}

	clone, err := CloneSemanticCorrection(correction)
	if err != nil {
		t.Fatalf("CloneSemanticCorrection() error = %v", err)
	}
	clone.AssistantResponse[0] = '['
	if got := string(correction.AssistantResponse); got != `{"items":["original"]}` {
		t.Fatalf("source correction changed through clone = %q", got)
	}
	if nilClone, err := CloneSemanticCorrection(nil); err != nil || nilClone != nil {
		t.Fatalf("CloneSemanticCorrection(nil) = %#v, %v; want nil, nil", nilClone, err)
	}

	encoded, err := json.Marshal(correction)
	if err != nil {
		t.Fatalf("marshal correction: %v", err)
	}
	if string(encoded) != "{}" {
		t.Fatalf("correction JSON = %s, want no sensitive content", encoded)
	}
}

func TestCorrectionContractsRejectInvalidContent(t *testing.T) {
	tooLongAssistant := bytes.Repeat([]byte("a"), MaxAssistantResponseBytes+1)
	tooLongGuidance := strings.Repeat("a", MaxCorrectionGuidanceBytes+1)
	tooLongReason := strings.Repeat("a", MaxValidationReasonCodeBytes+1)
	tooLongValidationGuidance := strings.Repeat("a", MaxValidationCorrectionGuidanceBytes+1)

	for _, test := range []struct {
		name string
		call func() error
	}{
		{"blank assistant", func() error { _, err := NewSemanticCorrection([]byte(" \n"), "guidance"); return err }},
		{"invalid assistant utf8", func() error { _, err := NewSemanticCorrection([]byte{0xff}, "guidance"); return err }},
		{"oversized assistant", func() error { _, err := NewSemanticCorrection(tooLongAssistant, "guidance"); return err }},
		{"blank guidance", func() error { _, err := NewSemanticCorrection([]byte("response"), " \t"); return err }},
		{"invalid guidance utf8", func() error { _, err := NewSemanticCorrection([]byte("response"), string([]byte{0xff})); return err }},
		{"oversized guidance", func() error { _, err := NewSemanticCorrection([]byte("response"), tooLongGuidance); return err }},
		{"unsupported protocol", func() error { _, err := NewModelCandidate([]byte("response"), "multiple_responses"); return err }},
		{"missing candidate protocol", func() error { _, err := NewModelCandidate([]byte("response"), ""); return err }},
		{"blank candidate response", func() error { _, err := NewModelCandidate([]byte(" "), CorrectionProtocolSingleResponseV1); return err }},
		{"missing rejection reason code", func() error {
			return ValidateValidationResult(ValidationResult{CorrectionGuidance: "Correct the response."})
		}},
		{"missing rejection guidance", func() error { return ValidateValidationResult(ValidationResult{ReasonCode: "invalid"}) }},
		{"oversized reason code", func() error {
			return ValidateValidationResult(ValidationResult{ReasonCode: tooLongReason, CorrectionGuidance: "Correct the response."})
		}},
		{"blank reason code", func() error {
			return ValidateValidationResult(ValidationResult{ReasonCode: " \t", CorrectionGuidance: "Correct the response."})
		}},
		{"invalid correction guidance utf8", func() error {
			return ValidateValidationResult(ValidationResult{ReasonCode: "invalid", CorrectionGuidance: string([]byte{0xff})})
		}},
		{"oversized correction guidance", func() error {
			return ValidateValidationResult(ValidationResult{ReasonCode: "invalid", CorrectionGuidance: tooLongValidationGuidance})
		}},
	} {
		t.Run(test.name, func(t *testing.T) {
			if err := test.call(); err == nil {
				t.Fatal("validation error = nil, want error")
			}
		})
	}
}

func TestModelCandidateOwnsValidatedResponse(t *testing.T) {
	response := []byte(`{"items":["original"]}`)
	candidate, err := NewModelCandidate(response, CorrectionProtocolSingleResponseV1)
	if err != nil {
		t.Fatalf("NewModelCandidate() error = %v", err)
	}
	response[0] = '['
	if got := string(candidate.Response); got != `{"items":["original"]}` {
		t.Fatalf("candidate response = %q, want owned original content", got)
	}
	clone, err := CloneModelCandidate(candidate)
	if err != nil {
		t.Fatalf("CloneModelCandidate() error = %v", err)
	}
	clone.Response[0] = '['
	if got := string(candidate.Response); got != `{"items":["original"]}` {
		t.Fatalf("source candidate changed through clone = %q", got)
	}
	if nilClone, err := CloneModelCandidate(nil); err != nil || nilClone != nil {
		t.Fatalf("CloneModelCandidate(nil) = %#v, %v; want nil, nil", nilClone, err)
	}
}

func TestValidationResultAllowsAbsentOptionalCorrectionFields(t *testing.T) {
	if err := ValidateValidationResult(ValidationResult{Approved: true}); err != nil {
		t.Fatalf("ValidateValidationResult() error = %v, want nil", err)
	}
	if err := ValidateValidationResult(ValidationResult{ReasonCode: "invalid-evidence", CorrectionGuidance: "Provide source-backed evidence."}); err != nil {
		t.Fatalf("ValidateValidationResult() error = %v, want nil", err)
	}
	if err := CorrectionProtocol("").Validate(); err == nil {
		t.Fatal("empty correction protocol validation error = nil, want error")
	}
}

func TestCloneStructuredCompletionRequestOwnsCorrection(t *testing.T) {
	correction, err := NewSemanticCorrection([]byte(`{"value":"original"}`), "Correct the value.")
	if err != nil {
		t.Fatalf("NewSemanticCorrection() error = %v", err)
	}
	attempts := 2
	request := StructuredCompletionRequest{
		Inputs:                         LLMInputSet{"source": NewLLMInputMaterial("source", "application/json", []byte(`{"source":true}`), "", "")},
		Vars:                           map[string]any{"labels": []string{"original"}},
		StructuredOutputRepairAttempts: &attempts,
		Correction:                     correction,
	}
	clone, err := CloneStructuredCompletionRequest(request)
	if err != nil {
		t.Fatalf("CloneStructuredCompletionRequest() error = %v", err)
	}
	correction.AssistantResponse[0] = '['
	request.Inputs["source"] = NewLLMInputMaterial("source", "application/json", []byte(`{"source":false}`), "", "")
	*request.StructuredOutputRepairAttempts = 7
	if got := string(clone.Correction.AssistantResponse); got != `{"value":"original"}` {
		t.Fatalf("cloned correction response = %q, want owned original content", got)
	}
	if got := string(clone.Inputs["source"].Content); got != `{"source":true}` {
		t.Fatalf("cloned input = %q, want owned original content", got)
	}
	if clone.StructuredOutputRepairAttempts == nil || *clone.StructuredOutputRepairAttempts != 2 {
		t.Fatalf("cloned repair attempts = %v, want 2", clone.StructuredOutputRepairAttempts)
	}
}

func TestStructuredCompletionRequestDebugSummaryOmitsCorrectionContent(t *testing.T) {
	const assistantResponse = `{"secret":"assistant response"}`
	const userGuidance = "secret user guidance"
	correction, err := NewSemanticCorrection([]byte(assistantResponse), userGuidance)
	if err != nil {
		t.Fatalf("NewSemanticCorrection() error = %v", err)
	}
	request := StructuredCompletionRequest{
		Inputs:     LLMInputSet{"source": NewLLMInputMaterial("source", "application/json", []byte(`{"source":true}`), "", "")},
		Vars:       map[string]any{"custom": "value"},
		Correction: correction,
	}

	summary := request.DebugSummary()
	if summary.InputCount != 1 || summary.VariableCount != 1 || summary.Correction == nil {
		t.Fatalf("debug summary = %#v, want input, variable, and correction metadata", summary)
	}
	if summary.Correction.AssistantResponseBytes != len(assistantResponse) || summary.Correction.UserGuidanceBytes != len(userGuidance) ||
		summary.Correction.AssistantResponseDigest == "" || summary.Correction.UserGuidanceDigest == "" {
		t.Fatalf("correction summary = %#v, want byte counts and digests", summary.Correction)
	}
	encoded, err := json.Marshal(summary)
	if err != nil {
		t.Fatalf("marshal debug summary: %v", err)
	}
	for _, rendered := range []string{string(encoded), fmt.Sprintf("%+v", request), fmt.Sprintf("%#v", request)} {
		for _, secret := range []string{assistantResponse, userGuidance} {
			if strings.Contains(rendered, secret) {
				t.Fatalf("content-safe request rendering leaked %q: %s", secret, rendered)
			}
		}
	}
}