diff --git a/internal/llm/adapter.go b/internal/llm/adapter.go index 0395d22..e6fa267 100644 --- a/internal/llm/adapter.go +++ b/internal/llm/adapter.go @@ -7,6 +7,7 @@ import ( "context" "errors" "fmt" + "io" "regexp" "strings" "unicode/utf8" @@ -206,7 +207,11 @@ func (e *StructuredValidationError) Error() string { if len(e.Attempts) > 1 && e.Attempts[1].DecodeError != nil { second = e.Attempts[1].DecodeError.Error() } - return fmt.Sprintf("%s: first: %s; second: %s", ErrStructuredOutputInvalidAfterRetry, first, second) + msg := fmt.Sprintf("%s: first: %s; second: %s", ErrStructuredOutputInvalidAfterRetry, first, second) + if len(e.Attempts) > 2 && e.Attempts[2].DecodeError != nil { + msg += fmt.Sprintf("; third: %s", e.Attempts[2].DecodeError.Error()) + } + return msg } // Is matches ErrStructuredOutputInvalidAfterRetry for errors.Is callers. @@ -255,16 +260,56 @@ func RunStructuredWithSessionResume[T any](ctx context.Context, adapter Adapter, return StructuredResult[T]{Response: retryResponse, SessionID: retrySessionID, ValidationAttempts: attempts}, err } retryValue, retryAcceptedOutput, retryErr := decodeStructuredAccepted(decode, retryResponse.StructuredOutput) - if retryErr != nil { + if retryErr == nil { + return StructuredResult[T]{Value: retryValue, Response: retryResponse, SessionID: retrySessionID, ValidationAttempts: attempts, AcceptedOutput: retryAcceptedOutput}, nil + } + attempts = append(attempts, StructuredValidationAttempt{ + Label: "retry", + SessionID: retrySessionID, + Response: cloneResponse(retryResponse), + DecodeError: retryErr, + }) + + // A truncated retry never finished, so the "invalid JSON" prompt can't fix + // it; give it one more attempt instead of failing the task. + if !isTruncatedStructuredOutput(retryErr) { + return StructuredResult[T]{Value: zero, Response: retryResponse, SessionID: retrySessionID, ValidationAttempts: attempts}, &StructuredValidationError{Attempts: attempts} + } + secondRetryReq := retryReq + if secondRetryReq.OnValidationRetry != nil { + if err := secondRetryReq.OnValidationRetry(&secondRetryReq); err != nil { + return StructuredResult[T]{Response: retryResponse, SessionID: retrySessionID, ValidationAttempts: attempts}, err + } + } + secondRetryReq.Prompt = retryPrompt(req.Prompt, retryErr) + secondRetryResumeSessionID := "" + if !secondRetryReq.FreshValidationRetrySession { + secondRetryResumeSessionID = retrySessionID + } + if strings.TrimSpace(secondRetryResumeSessionID) == "" && !secondRetryReq.FreshValidationRetrySession { + secondRetryResumeSessionID = resumeSessionID + } + secondRetrySessionID, secondRetryResponse, err := runOnceWithSession(ctx, adapter, secondRetryResumeSessionID, secondRetryReq) + if err != nil { + return StructuredResult[T]{Response: secondRetryResponse, SessionID: secondRetrySessionID, ValidationAttempts: attempts}, err + } + secondRetryValue, secondRetryAcceptedOutput, secondRetryErr := decodeStructuredAccepted(decode, secondRetryResponse.StructuredOutput) + if secondRetryErr != nil { attempts = append(attempts, StructuredValidationAttempt{ - Label: "retry", - SessionID: retrySessionID, - Response: cloneResponse(retryResponse), - DecodeError: retryErr, + Label: "retry_2", + SessionID: secondRetrySessionID, + Response: cloneResponse(secondRetryResponse), + DecodeError: secondRetryErr, }) - return StructuredResult[T]{Value: zero, Response: retryResponse, SessionID: retrySessionID, ValidationAttempts: attempts}, &StructuredValidationError{Attempts: attempts} + return StructuredResult[T]{Value: zero, Response: secondRetryResponse, SessionID: secondRetrySessionID, ValidationAttempts: attempts}, &StructuredValidationError{Attempts: attempts} } - return StructuredResult[T]{Value: retryValue, Response: retryResponse, SessionID: retrySessionID, ValidationAttempts: attempts, AcceptedOutput: retryAcceptedOutput}, nil + return StructuredResult[T]{Value: secondRetryValue, Response: secondRetryResponse, SessionID: secondRetrySessionID, ValidationAttempts: attempts, AcceptedOutput: secondRetryAcceptedOutput}, nil +} + +// isTruncatedStructuredOutput reports whether decoding stopped before a +// complete JSON value: io.EOF when empty, io.ErrUnexpectedEOF when cut off. +func isTruncatedStructuredOutput(err error) bool { + return errors.Is(err, io.ErrUnexpectedEOF) || errors.Is(err, io.EOF) } // decodeStructuredAccepted strict-decodes data, then on failure recovers a diff --git a/internal/llm/adapter_test.go b/internal/llm/adapter_test.go index 8e9d0fe..fd7a9a2 100644 --- a/internal/llm/adapter_test.go +++ b/internal/llm/adapter_test.go @@ -438,6 +438,98 @@ func assertValidationAttempts(t *testing.T, attempts []StructuredValidationAttem } } +func TestRunStructuredTruncatedResponseGetsExtraAttempt(t *testing.T) { + type probe struct { + OK bool `json:"ok"` + } + decodeProbe := func(data []byte) (probe, error) { + var p probe + if err := json.NewDecoder(bytes.NewReader(data)).Decode(&p); err != nil { + return probe{}, err + } + return p, nil + } + + t.Run("recovers on a third attempt after two truncated responses", func(t *testing.T) { + adapter := &FakeAdapter{} + adapter.Queue(FakeResult{SessionID: "s1", Response: Response{StructuredOutput: []byte(`{"ok":true,"ex`)}}) + adapter.Queue(FakeResult{SessionID: "s2", Response: Response{StructuredOutput: []byte(`{"ok":true,"ex`)}}) + adapter.Queue(FakeResult{SessionID: "s3", Response: Response{StructuredOutput: []byte(`{"ok":true}`)}}) + + result, err := RunStructuredWithSessionResume(context.Background(), adapter, "", Request{Prompt: "prompt"}, decodeProbe) + if err != nil { + t.Fatalf("RunStructuredWithSessionResume: %v", err) + } + if !result.Value.OK { + t.Fatalf("value = %#v, want recovered object", result.Value) + } + if requests := len(adapter.Requests()); requests != 3 { + t.Fatalf("requests = %d, want three attempts for a repeatedly truncated response", requests) + } + if len(result.ValidationAttempts) != 2 { + t.Fatalf("validation attempts = %#v, want two recorded failures before recovery", result.ValidationAttempts) + } + if result.ValidationAttempts[0].Label != "initial" || result.ValidationAttempts[1].Label != "retry" { + t.Fatalf("validation attempt labels = %#v, want initial then retry", result.ValidationAttempts) + } + if !errors.Is(result.ValidationAttempts[0].DecodeError, io.ErrUnexpectedEOF) || + !errors.Is(result.ValidationAttempts[1].DecodeError, io.ErrUnexpectedEOF) { + t.Fatalf("validation attempt decode errors = %#v, want unexpected EOF", result.ValidationAttempts) + } + }) + + t.Run("fails after three truncated attempts", func(t *testing.T) { + adapter := &FakeAdapter{} + adapter.Queue(FakeResult{SessionID: "s1", Response: Response{StructuredOutput: []byte(`{"ok":true,"ex`)}}) + adapter.Queue(FakeResult{SessionID: "s2", Response: Response{StructuredOutput: []byte(`{"ok":true,"ex`)}}) + adapter.Queue(FakeResult{SessionID: "s3", Response: Response{StructuredOutput: []byte(`{"ok":true,"ex`)}}) + + _, err := RunStructuredWithSessionResume(context.Background(), adapter, "", Request{Prompt: "prompt"}, decodeProbe) + if !errors.Is(err, ErrStructuredOutputInvalidAfterRetry) { + t.Fatalf("RunStructuredWithSessionResume error = %v, want %v", err, ErrStructuredOutputInvalidAfterRetry) + } + var validationErr *StructuredValidationError + if !errors.As(err, &validationErr) { + t.Fatalf("error type = %T, want StructuredValidationError", err) + } + if len(validationErr.Attempts) != 3 { + t.Fatalf("attempts = %#v, want three recorded attempts", validationErr.Attempts) + } + if validationErr.Attempts[2].Label != "retry_2" { + t.Fatalf("attempts[2].Label = %q, want retry_2", validationErr.Attempts[2].Label) + } + if !strings.Contains(err.Error(), "third:") { + t.Fatalf("error = %q, want a third summary once a third attempt ran", err.Error()) + } + if requests := len(adapter.Requests()); requests != 3 { + t.Fatalf("requests = %d, want exactly three attempts, not unbounded retry", requests) + } + }) + + t.Run("does not extend a retry that fails on a genuine schema violation", func(t *testing.T) { + adapter := &FakeAdapter{} + adapter.Queue(FakeResult{SessionID: "s1", Response: Response{StructuredOutput: []byte(`{"ok":true,"ex`)}}) + adapter.Queue(FakeResult{SessionID: "s2", Response: Response{StructuredOutput: []byte(`{"ok":false}`)}}) + + _, err := RunStructuredWithSessionResume(context.Background(), adapter, "", Request{Prompt: "prompt"}, func(data []byte) (probe, error) { + p, err := decodeProbe(data) + if err != nil { + return probe{}, err + } + if !p.OK { + return probe{}, errors.New("ok must be true") + } + return p, nil + }) + if !errors.Is(err, ErrStructuredOutputInvalidAfterRetry) { + t.Fatalf("RunStructuredWithSessionResume error = %v, want %v", err, ErrStructuredOutputInvalidAfterRetry) + } + if requests := len(adapter.Requests()); requests != 2 { + t.Fatalf("requests = %d, want the retry budget unchanged once the response is complete-but-wrong", requests) + } + }) +} + func TestRunStructuredProseRecovery(t *testing.T) { type probe struct { OK bool `json:"ok"` diff --git a/internal/pipeline/pipeline_test.go b/internal/pipeline/pipeline_test.go index af43d65..e5c92d7 100644 --- a/internal/pipeline/pipeline_test.go +++ b/internal/pipeline/pipeline_test.go @@ -7250,7 +7250,7 @@ func (a *reviewerIsolationAdapter) Start(_ context.Context, req llm.Request) (ll if attempt > 1 { sessionID = "beta-retry-session" } - return staticStream{sessionID: sessionID, output: `{"schema_version": 1, "agent_id": "harness:beta", "findings": [`}, nil + return staticStream{sessionID: sessionID, output: `{"schema_version": 1, "agent_id": "harness:beta", "findings": []}`}, nil case strings.Contains(req.Prompt, `"id": "harness:gamma"`): a.waitReviewerStart("harness:gamma") return staticStream{sessionID: "gamma-session", output: findingsJSON("harness:gamma", "main.go", "minor", 2, "gamma finding")}, nil