Build bounded output repair requests

docs/roadmap/implementation.md

@@ -333,6 +333,8 @@ Stage 3 is complete when the default repairer produces bounded, mode-correct,
 full-context requests without mutating prepared state or retaining failed
 history.
 
+**Status:** Complete.
+
 ## Stage 4: Complete The Internal Repair State Machine And Error Semantics
 
 ### Objective

internal/usecase/repairer.go

@@ -2,19 +2,29 @@ package usecase
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"strings"
+	"unicode/utf8"
 
 	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
 	"gitea.maximumdirect.net/eric/promptkit/internal/llm"
 )
 
+const (
+	maxRepairDiagnosticBytes = 64 * 1024
+	omittedRepairDiagnostics = "additional validation diagnostics were omitted"
+)
+
+// OutputRepairer generates a corrected candidate after validation fails.
 type OutputRepairer interface {
 	Repair(ctx context.Context, req RepairRequest) (*domain.GenerateResponse, error)
 }
 
+// RepairRequest contains the immutable execution state needed for one correction.
 type RepairRequest struct {
+	OriginalMessages []domain.RenderedMessage
 	PreviousOutput   string
 	ValidationErrors []string
 	SessionID        string
@@ -30,6 +40,7 @@ type defaultOutputRepairer struct {
 	llm llm.Client
 }
 
+// NewDefaultOutputRepairer constructs the standard internal output repairer.
 func NewDefaultOutputRepairer(llmClient llm.Client) OutputRepairer {
 	return &defaultOutputRepairer{llm: llmClient}
 }
@@ -39,33 +50,42 @@ func (r *defaultOutputRepairer) Repair(ctx context.Context, req RepairRequest) (
 		return nil, errors.New("llm client is required for repair")
 	}
 
-	errs := "(none provided)"
+	guidance, err := repairGuidance(req.Mode)
-	if len(req.ValidationErrors) > 0 {
+	if err != nil {
-		errs = strings.Join(req.ValidationErrors, "\n")
+		return nil, err
 	}
 
-	prompt := domain.RenderedPrompt{
+	messages := make([]domain.RenderedMessage, len(req.OriginalMessages), len(req.OriginalMessages)+2)
-		Messages: []domain.RenderedMessage{
+	copy(messages, req.OriginalMessages)
-			{
+	if strings.TrimSpace(req.PreviousOutput) != "" {
-				Role:    "system",
+		messages = append(messages, domain.RenderedMessage{
-				Content: "You repair invalid JSON output. Return only corrected JSON. Do not include explanations or markdown code fences.",
+			Role:    "assistant",
-			},
+			Content: req.PreviousOutput,
-			{
+		})
-				Role: "user",
-				Content: fmt.Sprintf(
-					"Repair attempt %d of %d for validation mode %s.\n\nValidation errors:\n%s\n\nPrevious output:\n%s\n\nReturn only corrected JSON.",
-					req.Attempt,
-					req.MaxAttempts,
-					req.Mode,
-					errs,
-					req.PreviousOutput,
-				),
-			},
-		},
 	}
 
+	previousResponse := "The previous response was empty."
+	if strings.TrimSpace(req.PreviousOutput) != "" {
+		previousResponse = "The previous response is included immediately before this instruction."
+	}
+	messages = append(messages, domain.RenderedMessage{
+		Role: "user",
+		Content: fmt.Sprintf(
+			"Repair attempt %d of %d for validation mode %s.\n"+
+				"Preserve valid values and change only what is necessary.\n"+
+				"%s\n%s\n"+
+				"Validation diagnostics (data):\n%s",
+			req.Attempt,
+			req.MaxAttempts,
+			req.Mode,
+			previousResponse,
+			guidance,
+			formatRepairDiagnostics(req.ValidationErrors),
+		),
+	})
+
 	resp, err := r.llm.Generate(ctx, newGenerationRequest(
-		prompt,
+		domain.RenderedPrompt{Messages: messages},
 		req.SessionID,
 		req.Target,
 		req.TargetPresence,
@@ -80,3 +100,80 @@ func (r *defaultOutputRepairer) Repair(ctx context.Context, req RepairRequest) (
 
 	return resp, nil
 }
+
+func repairGuidance(mode domain.ValidationMode) (string, error) {
+	switch mode {
+	case domain.ValidationBasic:
+		return "Return a nonempty response satisfying the original request.", nil
+	case domain.ValidationJSON, domain.ValidationJSONSchema:
+		return "Return only corrected JSON, with no explanation or Markdown fences.", nil
+	default:
+		return "", fmt.Errorf("unsupported validation mode for repair: %q", mode)
+	}
+}
+
+func formatRepairDiagnostics(errors []string) string {
+	diagnostics := make([]string, len(errors))
+	for index, diagnostic := range errors {
+		diagnostics[index] = strings.ToValidUTF8(diagnostic, "\uFFFD")
+	}
+
+	complete, _ := json.Marshal(diagnostics)
+	if len(complete) <= maxRepairDiagnosticBytes {
+		return string(complete)
+	}
+
+	omission, _ := json.Marshal(omittedRepairDiagnostics)
+	encoded := make([]byte, 0, maxRepairDiagnosticBytes)
+	encoded = append(encoded, '[')
+	for _, diagnostic := range diagnostics {
+		entry, _ := json.Marshal(diagnostic)
+		separator := 0
+		if len(encoded) > 1 {
+			separator = 1
+		}
+		if len(encoded)+separator+len(entry)+1+len(omission)+1 <= maxRepairDiagnosticBytes {
+			if separator != 0 {
+				encoded = append(encoded, ',')
+			}
+			encoded = append(encoded, entry...)
+			continue
+		}
+
+		available := maxRepairDiagnosticBytes - len(encoded) - separator - 1 - len(omission) - 1
+		if separator != 0 {
+			encoded = append(encoded, ',')
+		}
+		encoded = append(encoded, truncateDiagnosticJSONValue(diagnostic, available)...)
+		break
+	}
+	if len(encoded) > 1 {
+		encoded = append(encoded, ',')
+	}
+	encoded = append(encoded, omission...)
+	encoded = append(encoded, ']')
+	return string(encoded)
+}
+
+func truncateDiagnosticJSONValue(value string, maxBytes int) []byte {
+	if maxBytes < len(`""`) {
+		return []byte(`""`)
+	}
+
+	low, high := 0, len(value)
+	best := []byte(`""`)
+	for low <= high {
+		mid := low + (high-low)/2
+		for mid > 0 && mid < len(value) && !utf8.RuneStart(value[mid]) {
+			mid--
+		}
+		candidate, _ := json.Marshal(value[:mid])
+		if len(candidate) <= maxBytes {
+			best = candidate
+			low = mid + 1
+			continue
+		}
+		high = mid - 1
+	}
+	return best
+}

internal/usecase/repairer_test.go

@@ -0,0 +1,332 @@
+package usecase
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"reflect"
+	"strings"
+	"testing"
+	"unicode/utf8"
+
+	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
+)
+
+type recordingRepairClient struct {
+	requests []domain.GenerateRequest
+	response *domain.GenerateResponse
+	err      error
+}
+
+func (c *recordingRepairClient) Generate(_ context.Context, req domain.GenerateRequest) (*domain.GenerateResponse, error) {
+	c.requests = append(c.requests, req)
+	return c.response, c.err
+}
+
+func TestDefaultOutputRepairerBuildsFullContextRequest(t *testing.T) {
+	client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "corrected"}}
+	repairer := NewDefaultOutputRepairer(client)
+	original := []domain.RenderedMessage{
+		{Role: "system", Content: "Follow the task.", CacheControl: &domain.CacheControl{Type: domain.CacheControlEphemeral, TTL: "1h"}},
+		{Role: "user", Content: "Summarize the report."},
+	}
+	before := append([]domain.RenderedMessage(nil), original...)
+	previous := strings.Repeat("candidate ", 12_000)
+	target := domain.ExecutionTarget{BackendID: "backend", Endpoint: "https://provider.example/v1", Model: "model"}
+	presence := domain.ExecutionTargetPresence{Temperature: true, TopP: true}
+	structured := &domain.StructuredOutputSpec{}
+
+	response, err := repairer.Repair(context.Background(), RepairRequest{
+		OriginalMessages: original,
+		PreviousOutput:   previous,
+		ValidationErrors: []string{"invalid JSON"},
+		SessionID:        "session",
+		Target:           target,
+		TargetPresence:   presence,
+		StructuredOutput: structured,
+		Attempt:          1,
+		MaxAttempts:      3,
+		Mode:             domain.ValidationJSONSchema,
+	})
+	if err != nil {
+		t.Fatalf("repair: %v", err)
+	}
+	if response == nil || response.Content != "corrected" {
+		t.Fatalf("response = %+v", response)
+	}
+	if !reflect.DeepEqual(original, before) {
+		t.Fatalf("original messages changed: got %#v, want %#v", original, before)
+	}
+	if len(client.requests) != 1 {
+		t.Fatalf("generation requests = %d, want 1", len(client.requests))
+	}
+
+	request := client.requests[0]
+	if request.Prompt.SessionID != "session" || !reflect.DeepEqual(request.Target, target) || request.TargetPresence != presence || request.StructuredOutput != structured {
+		t.Fatalf("generation request fields = %+v", request)
+	}
+	if len(request.Prompt.Messages) != len(original)+2 {
+		t.Fatalf("message count = %d, want %d", len(request.Prompt.Messages), len(original)+2)
+	}
+	if !reflect.DeepEqual(request.Prompt.Messages[:len(original)], original) {
+		t.Fatalf("original messages = %#v, want %#v", request.Prompt.Messages[:len(original)], original)
+	}
+	assistant := request.Prompt.Messages[len(original)]
+	if assistant.Role != "assistant" || assistant.Content != previous {
+		t.Fatalf("assistant candidate = %+v", assistant)
+	}
+	correction := request.Prompt.Messages[len(original)+1]
+	if correction.Role != "user" || !strings.Contains(correction.Content, "Repair attempt 1 of 3") ||
+		!strings.Contains(correction.Content, "Preserve valid values") ||
+		!strings.Contains(correction.Content, "Return only corrected JSON") {
+		t.Fatalf("correction message = %q", correction.Content)
+	}
+	if diagnostics := repairDiagnosticsFromMessage(t, correction.Content); !reflect.DeepEqual(diagnostics, []string{"invalid JSON"}) {
+		t.Fatalf("diagnostics = %#v", diagnostics)
+	}
+}
+
+func TestDefaultOutputRepairerDoesNotAccumulateCandidates(t *testing.T) {
+	client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "corrected"}}
+	repairer := NewDefaultOutputRepairer(client)
+	backing := make([]domain.RenderedMessage, 1, 4)
+	backing[0] = domain.RenderedMessage{Role: "user", Content: "Original task"}
+	before := append([]domain.RenderedMessage(nil), backing...)
+
+	for _, candidate := range []string{"first invalid", "second invalid"} {
+		_, err := repairer.Repair(context.Background(), RepairRequest{
+			OriginalMessages: backing,
+			PreviousOutput:   candidate,
+			Attempt:          1,
+			MaxAttempts:      3,
+			Mode:             domain.ValidationJSON,
+		})
+		if err != nil {
+			t.Fatalf("repair %q: %v", candidate, err)
+		}
+	}
+	if !reflect.DeepEqual(backing, before) {
+		t.Fatalf("caller messages changed: got %#v, want %#v", backing, before)
+	}
+	if len(client.requests) != 2 {
+		t.Fatalf("generation requests = %d, want 2", len(client.requests))
+	}
+	for index, request := range client.requests {
+		messages := request.Prompt.Messages
+		if len(messages) != 3 || messages[0] != backing[0] || messages[1].Role != "assistant" || messages[1].Content != []string{"first invalid", "second invalid"}[index] {
+			t.Fatalf("request %d messages = %#v", index, messages)
+		}
+	}
+}
+
+func TestDefaultOutputRepairerOmitsEmptyCandidateMessage(t *testing.T) {
+	for _, tc := range []struct {
+		name      string
+		candidate string
+	}{
+		{name: "empty", candidate: ""},
+		{name: "whitespace", candidate: " \n\t "},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "corrected"}}
+			repairer := NewDefaultOutputRepairer(client)
+			_, err := repairer.Repair(context.Background(), RepairRequest{
+				OriginalMessages: []domain.RenderedMessage{{Role: "user", Content: "Original task"}},
+				PreviousOutput:   tc.candidate,
+				Attempt:          1,
+				MaxAttempts:      1,
+				Mode:             domain.ValidationBasic,
+			})
+			if err != nil {
+				t.Fatalf("repair: %v", err)
+			}
+			messages := client.requests[0].Prompt.Messages
+			if len(messages) != 2 || messages[1].Role != "user" || !strings.Contains(messages[1].Content, "previous response was empty") {
+				t.Fatalf("messages = %#v", messages)
+			}
+			if strings.Contains(messages[1].Content, tc.candidate) && tc.candidate != "" {
+				t.Fatalf("correction message repeated whitespace candidate: %q", messages[1].Content)
+			}
+		})
+	}
+}
+
+func TestDefaultOutputRepairerUsesModeSpecificGuidance(t *testing.T) {
+	tests := []struct {
+		mode domain.ValidationMode
+		want string
+	}{
+		{mode: domain.ValidationBasic, want: "Return a nonempty response"},
+		{mode: domain.ValidationJSON, want: "Return only corrected JSON"},
+		{mode: domain.ValidationJSONSchema, want: "Return only corrected JSON"},
+	}
+
+	for _, tc := range tests {
+		t.Run(string(tc.mode), func(t *testing.T) {
+			client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "corrected"}}
+			_, err := NewDefaultOutputRepairer(client).Repair(context.Background(), RepairRequest{Mode: tc.mode, Attempt: 1, MaxAttempts: 1})
+			if err != nil {
+				t.Fatalf("repair: %v", err)
+			}
+			if !strings.Contains(client.requests[0].Prompt.Messages[0].Content, tc.want) {
+				t.Fatalf("correction message = %q", client.requests[0].Prompt.Messages[0].Content)
+			}
+		})
+	}
+
+	client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "unexpected"}}
+	_, err := NewDefaultOutputRepairer(client).Repair(context.Background(), RepairRequest{Mode: domain.ValidationNone})
+	if err == nil || !strings.Contains(err.Error(), "unsupported validation mode") {
+		t.Fatalf("unsupported mode error = %v", err)
+	}
+	if len(client.requests) != 0 {
+		t.Fatalf("generation requests = %d, want 0", len(client.requests))
+	}
+}
+
+func TestFormatRepairDiagnosticsBoundsAndPreservesData(t *testing.T) {
+	tests := []struct {
+		name         string
+		errors       []string
+		wantOmission bool
+		check        func(*testing.T, []string)
+	}{
+		{
+			name:   "below limit",
+			errors: []string{"first", "second"},
+			check: func(t *testing.T, got []string) {
+				t.Helper()
+				if !reflect.DeepEqual(got, []string{"first", "second"}) {
+					t.Fatalf("diagnostics = %#v", got)
+				}
+			},
+		},
+		{
+			name:   "at limit",
+			errors: []string{strings.Repeat("x", maxRepairDiagnosticBytes-4)},
+			check: func(t *testing.T, got []string) {
+				t.Helper()
+				if len(got) != 1 || len(got[0]) != maxRepairDiagnosticBytes-4 {
+					t.Fatalf("diagnostics lengths = %#v", got)
+				}
+			},
+		},
+		{
+			name:         "multibyte truncation preserves prior entries",
+			errors:       []string{"first", strings.Repeat("界", maxRepairDiagnosticBytes)},
+			wantOmission: true,
+			check: func(t *testing.T, got []string) {
+				t.Helper()
+				if len(got) != 3 || got[0] != "first" || !strings.HasPrefix(strings.Repeat("界", maxRepairDiagnosticBytes), got[1]) {
+					t.Fatalf("diagnostics = %#v", got)
+				}
+			},
+		},
+		{
+			name:         "many diagnostics",
+			errors:       manyRepairDiagnostics(),
+			wantOmission: true,
+			check: func(t *testing.T, got []string) {
+				t.Helper()
+				if len(got) < 2 || got[0] != manyRepairDiagnostics()[0] {
+					t.Fatalf("diagnostics = %#v", got)
+				}
+			},
+		},
+		{
+			name:   "invalid UTF-8",
+			errors: []string{"broken\xffinput"},
+			check: func(t *testing.T, got []string) {
+				t.Helper()
+				if !reflect.DeepEqual(got, []string{"broken\uFFFDinput"}) {
+					t.Fatalf("diagnostics = %#v", got)
+				}
+			},
+		},
+		{
+			name:         "one huge diagnostic",
+			errors:       []string{strings.Repeat("x", maxRepairDiagnosticBytes*2)},
+			wantOmission: true,
+			check: func(t *testing.T, got []string) {
+				t.Helper()
+				if len(got) != 2 || !strings.HasPrefix(strings.Repeat("x", maxRepairDiagnosticBytes*2), got[0]) {
+					t.Fatalf("diagnostics = %#v", got)
+				}
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			before := append([]string(nil), tc.errors...)
+			encoded := formatRepairDiagnostics(tc.errors)
+			if len(encoded) > maxRepairDiagnosticBytes || !utf8.ValidString(encoded) {
+				t.Fatalf("encoded diagnostic length/UTF-8 = (%d, %v)", len(encoded), utf8.ValidString(encoded))
+			}
+			var got []string
+			if err := json.Unmarshal([]byte(encoded), &got); err != nil {
+				t.Fatalf("decode diagnostics: %v; encoded=%q", err, encoded)
+			}
+			if !reflect.DeepEqual(tc.errors, before) {
+				t.Fatalf("input diagnostics changed: got %#v, want %#v", tc.errors, before)
+			}
+			if hasOmission := len(got) > 0 && got[len(got)-1] == omittedRepairDiagnostics; hasOmission != tc.wantOmission {
+				t.Fatalf("omission = %v, want %v; diagnostics=%#v", hasOmission, tc.wantOmission, got)
+			}
+			tc.check(t, got)
+		})
+	}
+}
+
+func manyRepairDiagnostics() []string {
+	diagnostics := make([]string, 1_000)
+	for index := range diagnostics {
+		diagnostics[index] = fmt.Sprintf("diagnostic %04d %s", index, strings.Repeat("x", 128))
+	}
+	return diagnostics
+}
+
+func TestDefaultOutputRepairerPropagatesGenerationFailures(t *testing.T) {
+	expected := errors.New("generation failed")
+	for _, tc := range []struct {
+		name   string
+		client *recordingRepairClient
+		want   error
+	}{
+		{name: "nil client", want: nil},
+		{name: "generation error", client: &recordingRepairClient{err: expected}, want: expected},
+		{name: "nil response", client: &recordingRepairClient{}, want: nil},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			var repairer OutputRepairer
+			if tc.client == nil {
+				repairer = NewDefaultOutputRepairer(nil)
+			} else {
+				repairer = NewDefaultOutputRepairer(tc.client)
+			}
+			response, err := repairer.Repair(context.Background(), RepairRequest{Mode: domain.ValidationJSON})
+			if response != nil || err == nil {
+				t.Fatalf("response/error = (%+v, %v)", response, err)
+			}
+			if tc.want != nil && !errors.Is(err, tc.want) {
+				t.Fatalf("error = %v, want %v", err, tc.want)
+			}
+		})
+	}
+}
+
+func repairDiagnosticsFromMessage(t *testing.T, message string) []string {
+	t.Helper()
+	const marker = "Validation diagnostics (data):\n"
+	index := strings.Index(message, marker)
+	if index < 0 {
+		t.Fatalf("missing diagnostics marker in %q", message)
+	}
+	encoded := message[index+len(marker):]
+	var diagnostics []string
+	if err := json.Unmarshal([]byte(encoded), &diagnostics); err != nil {
+		t.Fatalf("decode diagnostics: %v", err)
+	}
+	return diagnostics
+}