From ddf73835593cd4481d18afad1f1029adc5f595e0 Mon Sep 17 00:00:00 2001
From: Lucas Machado <lucas.dpmachado@gmail.com>
Date: Sun, 8 Mar 2026 16:40:15 +0100
Subject: [PATCH] fix: recover truncated JSON responses and retry LLM
 generation

Add truncated JSON recovery that extracts release_notes from partial
responses when the LLM hits its token limit. Also adds retry logic
(up to 2 attempts) for both API failures and parse failures.
---
 actions/generate.go | 37 +++++++++++++++++------
 services/prompt.go  | 71 ++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 95 insertions(+), 13 deletions(-)

diff --git a/actions/generate.go b/actions/generate.go
index 698a0d7..0511e0a 100644
--- a/actions/generate.go
+++ b/actions/generate.go
@@ -240,16 +240,35 @@ func (a *GenerateAction) generateWithLLM(
 		existingTags,
 	)
 
-	helpers.Log.Info().Msgf("Sending prompt to %s (%s)...", inputs.Provider, inputs.Model)
-	start := time.Now()
+	maxRetries := 2
+	var lastErr error
 
-	response, err := a.llmSvc.Generate(inputs.Provider, inputs.Key, inputs.Model, prompt)
-	if err != nil {
-		return nil, err
-	}
+	for attempt := 1; attempt <= maxRetries; attempt++ {
+		if attempt > 1 {
+			helpers.Log.Info().Msgf("Retrying LLM generation (attempt %d/%d)...", attempt, maxRetries)
+		}
+
+		helpers.Log.Info().Msgf("Sending prompt to %s (%s)...", inputs.Provider, inputs.Model)
+		start := time.Now()
+
+		response, err := a.llmSvc.Generate(inputs.Provider, inputs.Key, inputs.Model, prompt)
+		if err != nil {
+			lastErr = err
+			continue
+		}
+
+		duration := time.Since(start).Seconds()
+		helpers.Log.Info().Msgf("LLM response received in %.2fs (%d characters)", duration, len(response))
+
+		result, err := a.promptSvc.ParseResponse(response)
+		if err != nil {
+			helpers.Log.Warn().Msgf("Parse failed on attempt %d: %v", attempt, err)
+			lastErr = err
+			continue
+		}
 
-	duration := time.Since(start).Seconds()
-	helpers.Log.Info().Msgf("LLM response received in %.2fs (%d characters)", duration, len(response))
+		return result, nil
+	}
 
-	return a.promptSvc.ParseResponse(response)
+	return nil, lastErr
 }
diff --git a/services/prompt.go b/services/prompt.go
index eb6d2be..ae2200b 100644
--- a/services/prompt.go
+++ b/services/prompt.go
@@ -172,18 +172,81 @@ func (p *PromptService) ParseResponse(text string) (*domain.StructuredResult, er
 		jsonEnd := strings.LastIndex(text, "}")
 		if jsonStart >= 0 && jsonEnd > jsonStart {
 			extracted := text[jsonStart : jsonEnd+1]
-			if err2 := json.Unmarshal([]byte(extracted), &result); err2 != nil {
-				return nil, fmt.Errorf("failed to parse LLM response as JSON: %w\nRaw output: %s", err, text[:min(500, len(text))])
+			if err2 := json.Unmarshal([]byte(extracted), &result); err2 == nil {
+				helpers.Log.Info().Msg("Structured JSON parsed successfully (extracted from text)")
+				return &result, nil
 			}
-		} else {
-			return nil, fmt.Errorf("failed to parse LLM response as JSON: %w\nRaw output: %s", err, text[:min(500, len(text))])
 		}
+
+		// Try to recover truncated JSON by extracting release_notes field
+		if recovered := recoverTruncatedJSON(text); recovered != nil {
+			helpers.Log.Warn().Msg("LLM response was truncated — recovered partial release notes")
+			return recovered, nil
+		}
+
+		return nil, fmt.Errorf("failed to parse LLM response as JSON: %w\nRaw output: %s", err, text[:min(500, len(text))])
 	}
 
 	helpers.Log.Info().Msg("Structured JSON parsed successfully")
 	return &result, nil
 }
 
+// recoverTruncatedJSON attempts to extract release_notes from a truncated JSON response.
+// This handles cases where the LLM hits its token limit mid-response.
+func recoverTruncatedJSON(text string) *domain.StructuredResult {
+	// Look for "release_notes" field value
+	re := regexp.MustCompile(`"release_notes"\s*:\s*"`)
+	loc := re.FindStringIndex(text)
+	if loc == nil {
+		return nil
+	}
+
+	// Extract the string value, handling escaped characters
+	start := loc[1]
+	var sb strings.Builder
+	i := start
+	for i < len(text) {
+		if text[i] == '\\' && i+1 < len(text) {
+			switch text[i+1] {
+			case '"':
+				sb.WriteByte('"')
+			case 'n':
+				sb.WriteByte('\n')
+			case 't':
+				sb.WriteByte('\t')
+			case '\\':
+				sb.WriteByte('\\')
+			default:
+				sb.WriteByte(text[i+1])
+			}
+			i += 2
+			continue
+		}
+		if text[i] == '"' {
+			break
+		}
+		sb.WriteByte(text[i])
+		i++
+	}
+
+	content := strings.TrimSpace(sb.String())
+	if content == "" {
+		return nil
+	}
+
+	// Try to extract suggested_version too
+	version := ""
+	versionRe := regexp.MustCompile(`"suggested_version"\s*:\s*"([^"]*)"`)
+	if m := versionRe.FindStringSubmatch(text); len(m) > 1 {
+		version = m[1]
+	}
+
+	return &domain.StructuredResult{
+		ReleaseNotes:     content,
+		SuggestedVersion: version,
+	}
+}
+
 func (p *PromptService) GenerateOutputPaths(basePath string) domain.OutputConfig {
 	if basePath == "" {
 		return domain.OutputConfig{