diff --git a/internal/usecase/repairer.go b/internal/usecase/repairer.go index aa83500..b0ed15c 100644 --- a/internal/usecase/repairer.go +++ b/internal/usecase/repairer.go @@ -132,7 +132,8 @@ func formatRepairDiagnostics(errors []string) string { if len(encoded) > 1 { separator = 1 } - if len(encoded)+separator+len(entry)+1+len(omission)+1 <= maxRepairDiagnosticBytes { + available := maxRepairDiagnosticBytes - len(encoded) - separator - 1 - len(omission) - 1 + if len(entry) <= available { if separator != 0 { encoded = append(encoded, ',') } @@ -140,7 +141,9 @@ func formatRepairDiagnostics(errors []string) string { continue } - available := maxRepairDiagnosticBytes - len(encoded) - separator - 1 - len(omission) - 1 + if available < len(`""`) { + break + } if separator != 0 { encoded = append(encoded, ',') } @@ -157,17 +160,24 @@ func formatRepairDiagnostics(errors []string) string { func truncateDiagnosticJSONValue(value string, maxBytes int) []byte { if maxBytes < len(`""`) { - return []byte(`""`) + return nil } - low, high := 0, len(value) + boundaries := []int{0} + for end := 0; end < len(value); { + _, size := utf8.DecodeRuneInString(value[end:]) + end += size + if end+len(`""`) > maxBytes { + break + } + boundaries = append(boundaries, end) + } + + low, high := 0, len(boundaries)-1 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]) + candidate, _ := json.Marshal(value[:boundaries[mid]]) if len(candidate) <= maxBytes { best = candidate low = mid + 1 diff --git a/internal/usecase/repairer_test.go b/internal/usecase/repairer_test.go index f575bf3..927fd12 100644 --- a/internal/usecase/repairer_test.go +++ b/internal/usecase/repairer_test.go @@ -234,6 +234,23 @@ func TestFormatRepairDiagnosticsBoundsAndPreservesData(t *testing.T) { } }, }, + { + name: "no room for partial diagnostic", + errors: func() []string { + omission, _ := json.Marshal(omittedRepairDiagnostics) + return []string{ + strings.Repeat("x", maxRepairDiagnosticBytes-len(omission)-5), + strings.Repeat("y", 128), + } + }(), + wantOmission: true, + check: func(t *testing.T, got []string) { + t.Helper() + if len(got) != 2 || got[1] != omittedRepairDiagnostics { + t.Fatalf("diagnostics = %#v", got) + } + }, + }, { name: "invalid UTF-8", errors: []string{"broken\xffinput"}, diff --git a/internal/usecase/runner_test.go b/internal/usecase/runner_test.go index c9c5dad..91e1736 100644 --- a/internal/usecase/runner_test.go +++ b/internal/usecase/runner_test.go @@ -295,8 +295,10 @@ func (c *controlledRepairLLM) Generate( ctx context.Context, req domain.GenerateRequest, ) (*domain.GenerateResponse, error) { - isRepair := len(req.Prompt.Messages) > 0 && - strings.HasPrefix(req.Prompt.Messages[0].Content, "You repair invalid JSON") + messages := req.Prompt.Messages + isRepair := len(messages) >= 2 && + messages[len(messages)-2].Role == "assistant" && + messages[len(messages)-1].Role == "user" c.mu.Lock() c.calls++ c.active++