Run bounded output repair attempts

docs/roadmap/implementation.md

@@ -430,6 +430,8 @@ Stage 4 is complete when the internal state machine satisfies every bounded
 success, exhaustion, error, prepared-state, and concurrency invariant without
 being publicly activated yet.
 
+**Status:** Complete.
+
 ## Stage 5: Activate Repair Through Public Engine Assembly
 
 ### Objective

internal/usecase/prepared_execution_test.go

@@ -3,10 +3,12 @@ package usecase
 import (
 	"context"
 	"errors"
+	"fmt"
 	"reflect"
 	"testing"
 
 	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
+	"gitea.maximumdirect.net/eric/promptkit/internal/llm"
 	"gitea.maximumdirect.net/eric/promptkit/internal/validate"
 )
 
@@ -385,19 +387,20 @@ func TestRunnerRunPreparedUsesFrozenValidationForInitialAndRepairOutputs(t *test
 	admitter := &fakeRunAdmitter{}
 	reader := defaultArtifactReader()
 	renderer := defaultRenderer()
+	client := &sequenceLLM{responses: []*domain.GenerateResponse{{
+		Content: `{"broken":true}`,
+		Usage: domain.TokenUsage{
+			PromptTokens: 13, CompletionTokens: 17, TotalTokens: 19,
+			CachedTokens: 23, CacheWriteTokens: 29,
+		},
+	}}}
 	runner := NewRunnerWithRepairer(
 		&fakePromptRepo{def: promptDef(domain.FormatJSON, domain.ValidationJSON, 1)},
 		&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}},
 		nil,
 		reader,
 		renderer,
-		&fakeLLM{resp: &domain.GenerateResponse{
+		client,
-			Content: `{"broken":true}`,
-			Usage: domain.TokenUsage{
-				PromptTokens: 13, CompletionTokens: 17, TotalTokens: 19,
-				CachedTokens: 23, CacheWriteTokens: 29,
-			},
-		}},
 		validator,
 		repairer,
 		admitter,
@@ -421,6 +424,10 @@ func TestRunnerRunPreparedUsesFrozenValidationForInitialAndRepairOutputs(t *test
 	if validator.directValidateCalls != 0 || repairer.calls != 1 {
 		t.Fatalf("validation/repair calls=(direct=%d repair=%d), want (0, 1)", validator.directValidateCalls, repairer.calls)
 	}
+	if len(client.requests) != 1 || len(repairer.reqs) != 1 ||
+		!reflect.DeepEqual(repairer.reqs[0].OriginalMessages, client.requests[0].Prompt.Messages) {
+		t.Fatalf("initial and repair messages = (%#v, %#v)", client.requests, repairer.reqs)
+	}
 	if result.Validation.Status != domain.ValidationPassed || result.Validation.RepairAttempts != 1 {
 		t.Fatalf("unexpected repaired validation result: %+v", result.Validation)
 	}
@@ -443,6 +450,7 @@ func TestRunnerRunPreparedReleasesAdmissionAcrossExecutionErrors(t *testing.T) {
 	generationFailure := errors.New("generation failed")
 	validationFailure := errors.New("validation failed")
 	repairFailure := errors.New("repair failed")
+	repairInvalidRequest := fmt.Errorf("repair request: %w", llm.ErrInvalidRequest)
 
 	tests := []struct {
 		name          string
@@ -450,6 +458,7 @@ func TestRunnerRunPreparedReleasesAdmissionAcrossExecutionErrors(t *testing.T) {
 		validation    *recordingPreparedValidation
 		repairer      *fakeRepairer
 		wantError     error
+		wantSource    error
 	}{
 		{
 			name:          "generation failure",
@@ -475,7 +484,50 @@ func TestRunnerRunPreparedReleasesAdmissionAcrossExecutionErrors(t *testing.T) {
 				}},
 			},
 			repairer:   &fakeRepairer{err: repairFailure},
-			wantError: ErrValidation,
+			wantError:  ErrLLMGenerate,
+			wantSource: repairFailure,
+		},
+		{
+			name: "repair invalid request",
+			validation: &recordingPreparedValidation{
+				results: []domain.ValidationResult{{
+					Status:  domain.ValidationFailed,
+					Mode:    domain.ValidationJSON,
+					Errors:  []string{"invalid"},
+					IsValid: false,
+				}},
+			},
+			repairer:   &fakeRepairer{err: repairInvalidRequest},
+			wantError:  ErrInvalidRequest,
+			wantSource: repairInvalidRequest,
+		},
+		{
+			name: "repair cancellation",
+			validation: &recordingPreparedValidation{
+				results: []domain.ValidationResult{{
+					Status:  domain.ValidationFailed,
+					Mode:    domain.ValidationJSON,
+					Errors:  []string{"invalid"},
+					IsValid: false,
+				}},
+			},
+			repairer:   &fakeRepairer{err: context.Canceled},
+			wantError:  ErrLLMGenerate,
+			wantSource: context.Canceled,
+		},
+		{
+			name: "repair deadline",
+			validation: &recordingPreparedValidation{
+				results: []domain.ValidationResult{{
+					Status:  domain.ValidationFailed,
+					Mode:    domain.ValidationJSON,
+					Errors:  []string{"invalid"},
+					IsValid: false,
+				}},
+			},
+			repairer:   &fakeRepairer{err: context.DeadlineExceeded},
+			wantError:  ErrLLMGenerate,
+			wantSource: context.DeadlineExceeded,
 		},
 	}
 
@@ -514,6 +566,9 @@ func TestRunnerRunPreparedReleasesAdmissionAcrossExecutionErrors(t *testing.T) {
 			if result != nil || !errors.Is(err, test.wantError) {
 				t.Fatalf("run prepared=(%+v, %v), want %v", result, err, test.wantError)
 			}
+			if test.wantSource != nil && !errors.Is(err, test.wantSource) {
+				t.Fatalf("run prepared error = %v, want source %v", err, test.wantSource)
+			}
 			if len(admitter.backendIDs) != 1 || admitter.releaseCalls != 1 {
 				t.Fatalf(
 					"admission calls=%#v releases=%d, want one each",

internal/usecase/runner.go

@@ -184,10 +184,10 @@ func (r *Runner) executePreparedRun(
 		prepared.StructuredOutput,
 	))
 	if err != nil {
-		if errors.Is(err, llm.ErrInvalidRequest) {
+		return nil, wrapGenerationError(err)
-			return nil, fmt.Errorf("%w: %w", ErrInvalidRequest, err)
 	}
-		return nil, fmt.Errorf("%w: %w", ErrLLMGenerate, err)
+	if genResp == nil {
+		return nil, fmt.Errorf("%w: model returned nil response", ErrLLMGenerate)
 	}
 	usage := genResp.Usage
 
@@ -203,6 +203,7 @@ func (r *Runner) executePreparedRun(
 			attemptsUsed++
 
 			repairResp, repairErr := r.repairer.Repair(ctx, RepairRequest{
+				OriginalMessages: prepared.Messages,
 				PreviousOutput:   genResp.Content,
 				ValidationErrors: validationResult.Errors,
 				SessionID:        prepared.SessionID,
@@ -214,10 +215,10 @@ func (r *Runner) executePreparedRun(
 				Mode:             prepared.OutputContract.ValidationMode,
 			})
 			if repairErr != nil {
-				return nil, fmt.Errorf("%w: %w", ErrValidation, repairErr)
+				return nil, wrapGenerationError(repairErr)
 			}
 			if repairResp == nil {
-				return nil, fmt.Errorf("%w: repairer returned nil response", ErrValidation)
+				return nil, fmt.Errorf("%w: repairer returned nil response", ErrLLMGenerate)
 			}
 
 			genResp = repairResp
@@ -257,6 +258,13 @@ func (r *Runner) executePreparedRun(
 	}, nil
 }
 
+func wrapGenerationError(err error) error {
+	if errors.Is(err, llm.ErrInvalidRequest) {
+		return fmt.Errorf("%w: %w", ErrInvalidRequest, err)
+	}
+	return fmt.Errorf("%w: %w", ErrLLMGenerate, err)
+}
+
 func addTokenUsage(total, next domain.TokenUsage) domain.TokenUsage {
 	return domain.TokenUsage{
 		PromptTokens:     total.PromptTokens + next.PromptTokens,
@@ -524,7 +532,12 @@ func (r *Runner) shouldAttemptRepair(contract domain.OutputContract, validationR
 	if validationResult.Status != domain.ValidationFailed {
 		return false
 	}
-	return contract.ValidationMode == domain.ValidationJSON || contract.ValidationMode == domain.ValidationJSONSchema
+	switch contract.ValidationMode {
+	case domain.ValidationBasic, domain.ValidationJSON, domain.ValidationJSONSchema:
+		return true
+	default:
+		return false
+	}
 }
 
 func mergeExecutionTarget(base domain.ExecutionTarget, override domain.ExecutionTarget) domain.ExecutionTarget {

internal/usecase/runner_test.go

@@ -2143,6 +2143,8 @@ func TestRunnerRepairStateMachine(t *testing.T) {
 		TopP:           true,
 		TimeoutSeconds: true,
 	}
+	emptyThenValid := responses(2)
+	emptyThenValid[0].Content = " \t "
 
 	tests := []struct {
 		name              string
@@ -2165,12 +2167,13 @@ func TestRunnerRepairStateMachine(t *testing.T) {
 			wantStatus:        domain.ValidationPassed,
 		},
 		{
-			name:              "basic failure is ineligible despite budget",
+			name:              "empty basic output repairs successfully",
 			mode:              domain.ValidationBasic,
 			budget:            3,
-			validationResults: []domain.ValidationResult{failed(domain.ValidationBasic, "empty output")},
+			validationResults: []domain.ValidationResult{failed(domain.ValidationBasic, "empty output"), passed(domain.ValidationBasic)},
-			responses:         responses(1),
+			responses:         emptyThenValid,
-			wantStatus:        domain.ValidationFailed,
+			wantRepairs:       1,
+			wantStatus:        domain.ValidationPassed,
 		},
 		{
 			name:   "inherited numeric values remain absent",
@@ -2305,6 +2308,9 @@ func TestRunnerRepairStateMachine(t *testing.T) {
 					!reflect.DeepEqual(req.ValidationErrors, tc.validationResults[index].Errors) {
 					t.Fatalf("repair request %d prior state = %+v", index, req)
 				}
+				if !reflect.DeepEqual(req.OriginalMessages, initialRequest.Prompt.Messages) {
+					t.Fatalf("repair request %d original messages drifted: %#v", index, req.OriginalMessages)
+				}
 				if req.TargetPresence != tc.wantPresence || !reflect.DeepEqual(req.Target, initialRequest.Target) ||
 					req.SessionID != initialRequest.Prompt.SessionID ||
 					!reflect.DeepEqual(req.StructuredOutput, initialRequest.StructuredOutput) {
@@ -2318,8 +2324,19 @@ func TestRunnerRepairStateMachine(t *testing.T) {
 					!reflect.DeepEqual(generated.StructuredOutput, initialRequest.StructuredOutput) {
 					t.Fatalf("repair generation request %d common fields drifted: %+v", index, generated)
 				}
-				if reflect.DeepEqual(generated.Prompt.Messages, initialRequest.Prompt.Messages) {
+				expectedMessages := len(initialRequest.Prompt.Messages) + 1
-					t.Fatalf("repair generation request %d reused the initial prompt", index)
+				if strings.TrimSpace(tc.responses[index].Content) != "" {
+					expectedMessages++
+				}
+				if len(generated.Prompt.Messages) != expectedMessages ||
+					generated.Prompt.Messages[len(generated.Prompt.Messages)-1].Role != "user" {
+					t.Fatalf("repair generation request %d messages = %#v", index, generated.Prompt.Messages)
+				}
+				if strings.TrimSpace(tc.responses[index].Content) != "" {
+					assistant := generated.Prompt.Messages[len(generated.Prompt.Messages)-2]
+					if assistant.Role != "assistant" || assistant.Content != tc.responses[index].Content {
+						t.Fatalf("repair generation request %d candidate = %+v", index, assistant)
+					}
 				}
 			}