518 lines
17 KiB
Go
518 lines
17 KiB
Go
package proposal_generation
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"os"
|
|
"path/filepath"
|
|
"reflect"
|
|
"runtime"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"gitea.maximumdirect.net/eric/audita/internal/core/config"
|
|
"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"
|
|
)
|
|
|
|
type fakeStructuredClient struct {
|
|
responses []StructuredCorrectionSet
|
|
err error
|
|
calls []contracts.StructuredCompletionRequest
|
|
}
|
|
|
|
func (f *fakeStructuredClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
|
|
_ = ctx
|
|
f.calls = append(f.calls, req)
|
|
if f.err != nil {
|
|
return contracts.StructuredCompletionResponse{}, f.err
|
|
}
|
|
target, ok := out.(*StructuredCorrectionSet)
|
|
if !ok {
|
|
return contracts.StructuredCompletionResponse{}, errors.New("unexpected output type")
|
|
}
|
|
if len(f.responses) == 0 {
|
|
return contracts.StructuredCompletionResponse{}, errors.New("unexpected call")
|
|
}
|
|
*target = f.responses[0]
|
|
f.responses = f.responses[1:]
|
|
return contracts.StructuredCompletionResponse{}, nil
|
|
}
|
|
|
|
type countingScheduler struct {
|
|
runs int
|
|
}
|
|
|
|
func (s *countingScheduler) Run(ctx context.Context, fn func(context.Context) error) error {
|
|
s.runs++
|
|
return fn(ctx)
|
|
}
|
|
|
|
type captureDiagnosticsWriter struct {
|
|
lastStage string
|
|
lastRequestMetadata any
|
|
lastRequestPayload any
|
|
}
|
|
|
|
func (w *captureDiagnosticsWriter) WriteInteraction(stage string, requestMetadata any, requestPayload any, responsePayload any, errorPayload any) (InteractionArtifacts, error) {
|
|
w.lastStage = stage
|
|
w.lastRequestMetadata = requestMetadata
|
|
w.lastRequestPayload = requestPayload
|
|
_ = responsePayload
|
|
_ = errorPayload
|
|
return InteractionArtifacts{}, nil
|
|
}
|
|
|
|
type sleepingStructuredClient struct {
|
|
inFlight int32
|
|
maxInFlight int32
|
|
entered chan struct{}
|
|
release chan struct{}
|
|
}
|
|
|
|
func (c *sleepingStructuredClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
|
|
_ = req
|
|
current := atomic.AddInt32(&c.inFlight, 1)
|
|
for {
|
|
prior := atomic.LoadInt32(&c.maxInFlight)
|
|
if current <= prior || atomic.CompareAndSwapInt32(&c.maxInFlight, prior, current) {
|
|
break
|
|
}
|
|
}
|
|
if c.entered != nil {
|
|
c.entered <- struct{}{}
|
|
}
|
|
select {
|
|
case <-c.release:
|
|
case <-ctx.Done():
|
|
atomic.AddInt32(&c.inFlight, -1)
|
|
return contracts.StructuredCompletionResponse{}, ctx.Err()
|
|
}
|
|
atomic.AddInt32(&c.inFlight, -1)
|
|
|
|
target, ok := out.(*StructuredCorrectionSet)
|
|
if !ok {
|
|
return contracts.StructuredCompletionResponse{}, errors.New("unexpected output type")
|
|
}
|
|
*target = StructuredCorrectionSet{
|
|
Corrections: []StructuredCorrectionProposal{
|
|
{TargetSegmentID: 1, OriginalText: "x", CorrectedText: "y", Confidence: 0.9},
|
|
},
|
|
}
|
|
return contracts.StructuredCompletionResponse{}, nil
|
|
}
|
|
|
|
func defaultRequest(t *testing.T) Request {
|
|
t.Helper()
|
|
cfg := config.Default()
|
|
return Request{
|
|
ModuleKey: "test_module",
|
|
ModuleInstance: "test_module",
|
|
ReplacementPolicy: proposals.ReplacementPolicyRequireUnique,
|
|
Config: &cfg,
|
|
Messages: []contracts.LLMMessage{
|
|
{Role: "system", Content: "system prompt"},
|
|
{Role: "user", Content: "user prompt"},
|
|
},
|
|
StartIndex: 0,
|
|
}
|
|
}
|
|
|
|
func TestGenerateCandidatesSuccess(t *testing.T) {
|
|
client := &fakeStructuredClient{
|
|
responses: []StructuredCorrectionSet{
|
|
{
|
|
Corrections: []StructuredCorrectionProposal{
|
|
{TargetSegmentID: 7, OriginalText: "teh", CorrectedText: "the", Confidence: 0.9},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
section := contracts.SectionMetadata{Index: 2}
|
|
req := defaultRequest(t)
|
|
req.LLMClient = client
|
|
req.StartIndex = 10
|
|
req.Section = §ion
|
|
|
|
got, err := GenerateCandidates(context.Background(), req)
|
|
if err != nil {
|
|
t.Fatalf("GenerateCandidates error: %v", err)
|
|
}
|
|
if len(got.Corrections) != 1 || len(got.Enriched) != 1 {
|
|
t.Fatalf("unexpected proposal lengths: %+v", got)
|
|
}
|
|
if got.Enriched[0].ProposalIndex != 10 {
|
|
t.Fatalf("expected proposal index 10, got %d", got.Enriched[0].ProposalIndex)
|
|
}
|
|
if got.Enriched[0].SectionIndex == nil || *got.Enriched[0].SectionIndex != 2 {
|
|
t.Fatalf("expected section index 2, got %v", got.Enriched[0].SectionIndex)
|
|
}
|
|
}
|
|
|
|
func TestGenerateCandidatesUsesCorrectionSetSchema(t *testing.T) {
|
|
client := &fakeStructuredClient{
|
|
responses: []StructuredCorrectionSet{
|
|
{Corrections: []StructuredCorrectionProposal{{TargetSegmentID: 1, OriginalText: "teh", CorrectedText: "the", Confidence: 0.9}}},
|
|
},
|
|
}
|
|
req := defaultRequest(t)
|
|
req.LLMClient = client
|
|
|
|
_, err := GenerateCandidates(context.Background(), req)
|
|
if err != nil {
|
|
t.Fatalf("GenerateCandidates error: %v", err)
|
|
}
|
|
if len(client.calls) != 1 {
|
|
t.Fatalf("expected 1 LLM call, got %d", len(client.calls))
|
|
}
|
|
call := client.calls[0]
|
|
if call.ResponseSchema == nil {
|
|
t.Fatalf("expected response schema on structured request")
|
|
}
|
|
want := responseschema.MustLookup(responseschema.CorrectionSetKey)
|
|
if call.ResponseSchema.ID != want.ID || call.ResponseSchema.Version != want.Version || call.ResponseSchema.Name != want.Name || call.ResponseSchema.SHA256 != want.SHA256 {
|
|
t.Fatalf("unexpected response schema metadata: got=%+v want=%+v", *call.ResponseSchema, want)
|
|
}
|
|
}
|
|
|
|
func TestGenerateCandidatesDiagnosticsIncludeSchemaMetadata(t *testing.T) {
|
|
client := &fakeStructuredClient{
|
|
responses: []StructuredCorrectionSet{
|
|
{Corrections: []StructuredCorrectionProposal{{TargetSegmentID: 1, OriginalText: "teh", CorrectedText: "the", Confidence: 0.9}}},
|
|
},
|
|
}
|
|
diag := &captureDiagnosticsWriter{}
|
|
req := defaultRequest(t)
|
|
req.LLMClient = client
|
|
req.DiagnosticsWriter = diag
|
|
req.PromptMetadata = map[string]any{
|
|
"prompt_id": "modules.grammar.proposal",
|
|
"prompt_version": "v1",
|
|
"prompt_source": "builtin",
|
|
"embedded_path": "assets/modules/grammar",
|
|
"sha256": "abc",
|
|
}
|
|
|
|
_, err := GenerateCandidates(context.Background(), req)
|
|
if err != nil {
|
|
t.Fatalf("GenerateCandidates error: %v", err)
|
|
}
|
|
metadata, ok := diag.lastRequestMetadata.(map[string]any)
|
|
if !ok {
|
|
t.Fatalf("expected request metadata map, got %T", diag.lastRequestMetadata)
|
|
}
|
|
schemaMap, ok := metadata["response_schema"].(map[string]any)
|
|
if !ok {
|
|
t.Fatalf("expected response_schema metadata map, got %T", metadata["response_schema"])
|
|
}
|
|
want := responseschema.MustLookup(responseschema.CorrectionSetKey)
|
|
if schemaMap["id"] != want.ID || schemaMap["version"] != want.Version || schemaMap["name"] != want.Name || schemaMap["sha256"] != want.SHA256 {
|
|
t.Fatalf("unexpected diagnostics schema metadata: got=%v want=%+v", schemaMap, want)
|
|
}
|
|
promptMap, ok := metadata["prompt_metadata"].(map[string]any)
|
|
if !ok {
|
|
t.Fatalf("expected prompt_metadata map, got %T", metadata["prompt_metadata"])
|
|
}
|
|
if promptMap["prompt_id"] != "modules.grammar.proposal" || promptMap["prompt_version"] != "v1" || promptMap["prompt_source"] != "builtin" || promptMap["embedded_path"] != "assets/modules/grammar" || promptMap["sha256"] != "abc" {
|
|
t.Fatalf("unexpected prompt metadata map: %v", promptMap)
|
|
}
|
|
}
|
|
|
|
func TestGenerateCandidatesInvalidCorrectionIsPreservedForLaterValidation(t *testing.T) {
|
|
client := &fakeStructuredClient{
|
|
responses: []StructuredCorrectionSet{
|
|
{
|
|
Corrections: []StructuredCorrectionProposal{
|
|
{TargetSegmentID: 0, OriginalText: "x", CorrectedText: "", Confidence: 1.2},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
req := defaultRequest(t)
|
|
req.LLMClient = client
|
|
|
|
result, err := GenerateCandidates(context.Background(), req)
|
|
if err != nil {
|
|
t.Fatalf("expected invalid correction to survive generation, got %v", err)
|
|
}
|
|
if len(result.Corrections) != 1 {
|
|
t.Fatalf("expected one correction, got %+v", result)
|
|
}
|
|
if result.Corrections[0].TargetSegmentID != 0 || result.Corrections[0].Confidence != 1.2 {
|
|
t.Fatalf("unexpected preserved correction: %+v", result.Corrections[0])
|
|
}
|
|
if len(result.Warnings) != 0 {
|
|
t.Fatalf("did not expect warnings for individually invalid corrections, got %+v", result.Warnings)
|
|
}
|
|
}
|
|
|
|
func TestGenerateCandidatesMalformedStructuredOutputReturnsWarning(t *testing.T) {
|
|
client := &fakeStructuredClient{err: errors.New("malformed structured output")}
|
|
req := defaultRequest(t)
|
|
req.LLMClient = client
|
|
|
|
result, err := GenerateCandidates(context.Background(), req)
|
|
if err != nil {
|
|
t.Fatalf("expected malformed structured output to downgrade to warning, got %v", err)
|
|
}
|
|
if len(result.Corrections) != 0 || len(result.Enriched) != 0 {
|
|
t.Fatalf("expected no proposals on malformed response, got %+v", result)
|
|
}
|
|
if len(result.Warnings) != 1 {
|
|
t.Fatalf("expected one warning, got %+v", result.Warnings)
|
|
}
|
|
if result.Warnings[0].ReasonCode != "proposal_response_malformed" {
|
|
t.Fatalf("unexpected warning: %+v", result.Warnings[0])
|
|
}
|
|
}
|
|
|
|
func TestGenerateCandidatesProviderMalformedEnvelopeReturnsWarning(t *testing.T) {
|
|
client := &fakeStructuredClient{err: errors.New("provider response missing choices")}
|
|
req := defaultRequest(t)
|
|
req.LLMClient = client
|
|
|
|
result, err := GenerateCandidates(context.Background(), req)
|
|
if err != nil {
|
|
t.Fatalf("expected malformed provider envelope to downgrade to warning, got %v", err)
|
|
}
|
|
if len(result.Corrections) != 0 || len(result.Enriched) != 0 {
|
|
t.Fatalf("expected no proposals on malformed response, got %+v", result)
|
|
}
|
|
if len(result.Warnings) != 1 || result.Warnings[0].ReasonCode != "proposal_response_malformed" {
|
|
t.Fatalf("unexpected warnings: %+v", result.Warnings)
|
|
}
|
|
}
|
|
|
|
func TestGenerateCandidatesDeterministicIndexAssignment(t *testing.T) {
|
|
baseResponse := StructuredCorrectionSet{
|
|
Corrections: []StructuredCorrectionProposal{
|
|
{TargetSegmentID: 1, OriginalText: "a", CorrectedText: "A", Confidence: 0.9},
|
|
{TargetSegmentID: 2, OriginalText: "b", CorrectedText: "B", Confidence: 0.9},
|
|
},
|
|
}
|
|
clientA := &fakeStructuredClient{responses: []StructuredCorrectionSet{baseResponse}}
|
|
clientB := &fakeStructuredClient{responses: []StructuredCorrectionSet{baseResponse}}
|
|
|
|
reqA := defaultRequest(t)
|
|
reqA.LLMClient = clientA
|
|
reqA.StartIndex = 3
|
|
first, err := GenerateCandidates(context.Background(), reqA)
|
|
if err != nil {
|
|
t.Fatalf("first generation failed: %v", err)
|
|
}
|
|
|
|
reqB := defaultRequest(t)
|
|
reqB.LLMClient = clientB
|
|
reqB.StartIndex = 3
|
|
second, err := GenerateCandidates(context.Background(), reqB)
|
|
if err != nil {
|
|
t.Fatalf("second generation failed: %v", err)
|
|
}
|
|
|
|
if !reflect.DeepEqual(first.Enriched, second.Enriched) {
|
|
t.Fatalf("expected stable enriched proposals\nfirst=%+v\nsecond=%+v", first.Enriched, second.Enriched)
|
|
}
|
|
}
|
|
|
|
func TestGenerateCandidatesMultipleSectionsStableMetadata(t *testing.T) {
|
|
client := &fakeStructuredClient{
|
|
responses: []StructuredCorrectionSet{
|
|
{Corrections: []StructuredCorrectionProposal{{TargetSegmentID: 1, OriginalText: "alpha", CorrectedText: "ALPHA", Confidence: 0.9}}},
|
|
{Corrections: []StructuredCorrectionProposal{{TargetSegmentID: 3, OriginalText: "charlie", CorrectedText: "CHARLIE", Confidence: 0.9}}},
|
|
},
|
|
}
|
|
section0 := contracts.SectionMetadata{Index: 0}
|
|
section1 := contracts.SectionMetadata{Index: 1}
|
|
|
|
req0 := defaultRequest(t)
|
|
req0.LLMClient = client
|
|
req0.Section = §ion0
|
|
req0.StartIndex = 0
|
|
part0, err := GenerateCandidates(context.Background(), req0)
|
|
if err != nil {
|
|
t.Fatalf("section 0 generation failed: %v", err)
|
|
}
|
|
|
|
req1 := defaultRequest(t)
|
|
req1.LLMClient = client
|
|
req1.Section = §ion1
|
|
req1.StartIndex = len(part0.Enriched)
|
|
part1, err := GenerateCandidates(context.Background(), req1)
|
|
if err != nil {
|
|
t.Fatalf("section 1 generation failed: %v", err)
|
|
}
|
|
|
|
all := append(append([]proposals.EnrichedCorrectionProposal(nil), part0.Enriched...), part1.Enriched...)
|
|
if len(all) != 2 {
|
|
t.Fatalf("expected 2 proposals, got %d", len(all))
|
|
}
|
|
if all[0].ProposalIndex != 0 || all[1].ProposalIndex != 1 {
|
|
t.Fatalf("unexpected proposal indexes: %d, %d", all[0].ProposalIndex, all[1].ProposalIndex)
|
|
}
|
|
if all[0].SectionIndex == nil || *all[0].SectionIndex != 0 || all[1].SectionIndex == nil || *all[1].SectionIndex != 1 {
|
|
t.Fatalf("unexpected section metadata: %+v", all)
|
|
}
|
|
}
|
|
|
|
func TestGenerateCandidatesDiagnosticsWrittenAndRedacted(t *testing.T) {
|
|
primarySecret := "proposal-primary-secret"
|
|
validationSecret := "proposal-validation-secret"
|
|
client := &fakeStructuredClient{
|
|
responses: []StructuredCorrectionSet{
|
|
{Corrections: []StructuredCorrectionProposal{{TargetSegmentID: 1, OriginalText: validationSecret, CorrectedText: "safe", Confidence: 0.9}}},
|
|
},
|
|
}
|
|
cfg := config.Default()
|
|
cfg.PrimaryLLM.APIKey = primarySecret
|
|
cfg.ValidationLLM.APIKey = validationSecret
|
|
req := defaultRequest(t)
|
|
req.Config = &cfg
|
|
req.LLMClient = client
|
|
req.DiagnosticsDir = t.TempDir()
|
|
req.Messages = []contracts.LLMMessage{
|
|
{Role: "system", Content: "include secret " + primarySecret},
|
|
{Role: "user", Content: "fix it"},
|
|
}
|
|
|
|
got, err := GenerateCandidates(context.Background(), req)
|
|
if err != nil {
|
|
t.Fatalf("GenerateCandidates error: %v", err)
|
|
}
|
|
if got.Artifacts.ResponsePayloadPath == "" || got.Artifacts.RequestPayloadPath == "" {
|
|
t.Fatalf("expected diagnostics artifact paths, got %+v", got.Artifacts)
|
|
}
|
|
|
|
for _, path := range []string{got.Artifacts.RequestPayloadPath, got.Artifacts.ResponsePayloadPath} {
|
|
raw, readErr := os.ReadFile(path)
|
|
if readErr != nil {
|
|
t.Fatalf("read artifact %q: %v", path, readErr)
|
|
}
|
|
if strings.Contains(string(raw), primarySecret) || strings.Contains(string(raw), validationSecret) {
|
|
t.Fatalf("artifact leaked configured secret in %q: %s", path, string(raw))
|
|
}
|
|
if !strings.Contains(string(raw), "[REDACTED]") {
|
|
t.Fatalf("expected redaction marker in artifact %q: %s", path, string(raw))
|
|
}
|
|
}
|
|
|
|
raw, readErr := os.ReadFile(got.Artifacts.RequestMetadataPath)
|
|
if readErr != nil {
|
|
t.Fatalf("read metadata artifact %q: %v", got.Artifacts.RequestMetadataPath, readErr)
|
|
}
|
|
var metadata map[string]any
|
|
if err := json.Unmarshal(raw, &metadata); err != nil {
|
|
t.Fatalf("unmarshal metadata artifact: %v", err)
|
|
}
|
|
schemaMap, ok := metadata["response_schema"].(map[string]any)
|
|
if !ok {
|
|
t.Fatalf("expected response_schema metadata in diagnostics, got %T", metadata["response_schema"])
|
|
}
|
|
want := responseschema.MustLookup(responseschema.CorrectionSetKey)
|
|
if schemaMap["id"] != want.ID || schemaMap["version"] != want.Version || schemaMap["name"] != want.Name || schemaMap["sha256"] != want.SHA256 {
|
|
t.Fatalf("unexpected schema metadata in diagnostics: got=%v want=%+v", schemaMap, want)
|
|
}
|
|
}
|
|
|
|
func TestGenerateCandidatesSchedulerUsage(t *testing.T) {
|
|
client := &fakeStructuredClient{
|
|
responses: []StructuredCorrectionSet{
|
|
{Corrections: []StructuredCorrectionProposal{{TargetSegmentID: 1, OriginalText: "x", CorrectedText: "y", Confidence: 0.9}}},
|
|
},
|
|
}
|
|
scheduler := &countingScheduler{}
|
|
req := defaultRequest(t)
|
|
req.LLMClient = client
|
|
req.Scheduler = scheduler
|
|
|
|
_, err := GenerateCandidates(context.Background(), req)
|
|
if err != nil {
|
|
t.Fatalf("GenerateCandidates error: %v", err)
|
|
}
|
|
if scheduler.runs != 1 {
|
|
t.Fatalf("expected scheduler to run once, got %d", scheduler.runs)
|
|
}
|
|
}
|
|
|
|
func TestGenerateCandidatesClientError(t *testing.T) {
|
|
client := &fakeStructuredClient{err: errors.New("boom")}
|
|
req := defaultRequest(t)
|
|
req.LLMClient = client
|
|
req.DiagnosticsDir = t.TempDir()
|
|
|
|
got, err := GenerateCandidates(context.Background(), req)
|
|
if err == nil || !strings.Contains(err.Error(), "completion failed") {
|
|
t.Fatalf("expected completion failure, got %v", err)
|
|
}
|
|
if got.Artifacts.ResponsePayloadPath != "" {
|
|
t.Fatalf("expected zero result on error, got %+v", got)
|
|
}
|
|
matches, globErr := filepath.Glob(filepath.Join(req.DiagnosticsDir, req.ModuleInstance, "*error-payload.json"))
|
|
if globErr != nil {
|
|
t.Fatalf("glob error: %v", globErr)
|
|
}
|
|
if len(matches) == 0 {
|
|
t.Fatalf("expected error diagnostics artifact under %s", req.DiagnosticsDir)
|
|
}
|
|
}
|
|
|
|
func TestGenerateCandidatesRespectsSchedulerConcurrency(t *testing.T) {
|
|
scheduler, err := llm.NewScheduler(2)
|
|
if err != nil {
|
|
t.Fatalf("NewScheduler: %v", err)
|
|
}
|
|
client := &sleepingStructuredClient{
|
|
entered: make(chan struct{}, 16),
|
|
release: make(chan struct{}),
|
|
}
|
|
baseReq := defaultRequest(t)
|
|
baseReq.LLMClient = client
|
|
baseReq.Scheduler = scheduler
|
|
|
|
var wg sync.WaitGroup
|
|
for i := 0; i < 8; i++ {
|
|
wg.Add(1)
|
|
go func(idx int) {
|
|
defer wg.Done()
|
|
req := baseReq
|
|
req.StartIndex = idx
|
|
if _, runErr := GenerateCandidates(context.Background(), req); runErr != nil {
|
|
t.Errorf("GenerateCandidates[%d] error: %v", idx, runErr)
|
|
}
|
|
}(i)
|
|
}
|
|
|
|
waitForEntries(t, client.entered, 2)
|
|
close(client.release)
|
|
wg.Wait()
|
|
|
|
if got := atomic.LoadInt32(&client.maxInFlight); got > 2 {
|
|
t.Fatalf("expected scheduler cap <= 2, got %d", got)
|
|
}
|
|
if got := atomic.LoadInt32(&client.maxInFlight); got < 2 {
|
|
t.Fatalf("expected observed concurrency of at least 2, got %d", got)
|
|
}
|
|
}
|
|
|
|
func waitForEntries(t *testing.T, entered <-chan struct{}, want int) {
|
|
t.Helper()
|
|
deadline := time.Now().Add(300 * time.Millisecond)
|
|
got := 0
|
|
for got < want && time.Now().Before(deadline) {
|
|
select {
|
|
case <-entered:
|
|
got++
|
|
default:
|
|
runtime.Gosched()
|
|
}
|
|
}
|
|
if got < want {
|
|
t.Fatalf("timed out waiting for %d entered calls (got %d)", want, got)
|
|
}
|
|
}
|