From a0b6a2a0e1807e299390aedee83b2393b24f5445 Mon Sep 17 00:00:00 2001
From: David Hoenisch <dh1689@pm.me>
Date: Tue, 28 Apr 2026 15:03:56 -0700
Subject: [PATCH] fix(ai): handle custom LLM models that emit test-case
 declarations

When using a custom Ollama model (e.g. llama3.2:latest) the LLM often
outputs the generated test cases as a Go variable declaration instead
of a bare array literal or a complete test function:

    tests := []struct {
        name string
        args args
        want int
    }{
        {name: "case 1", args: args{...}, want: 42},
    }

Two parser bugs caused these responses to be rejected:

1. parseGoTestCases() only recognised bare composite literals or
   complete func Test... blocks. The tests := form fell through to
   parseTestCaseArray, which wrapped it in another literal producing
   invalid Go.

2. ensureTrailingComma() saw the declaration ended with '}' and was
   not a func Test, so it appended a comma: "tests := []struct{}{...},"
   which is invalid Go syntax.

Both issues caused GenerateTestCases() to return an error. The error
output (which contained the LLM's generated cases) was printed to
stderr/stdout, while the actual test file only received the TODO
fallback comment.

Fix:
- Detect tests := / tests= declarations in parseGoTestCases and wrap
  them in a func init() so the existing AST parser can extract cases.
- Teach ensureTrailingComma to skip tests := and var tests declarations.

Verified with llama3.2:latest and qwen2.5-coder:0.5b.
---
 internal/ai/parser_go.go | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/internal/ai/parser_go.go b/internal/ai/parser_go.go
index 4d35a6d5..371b79ee 100644
--- a/internal/ai/parser_go.go
+++ b/internal/ai/parser_go.go
@@ -35,12 +35,14 @@ func extractCodeFromMarkdown(text string) string {
 }
 
 // ensureTrailingComma adds a trailing comma to the last struct if missing.
-// Only applies when parsing test case arrays, not complete functions.
+// Only applies when parsing test case arrays, not complete functions or declarations.
 func ensureTrailingComma(code string) string {
 	code = strings.TrimSpace(code)
 
-	// Don't modify complete functions
-	if strings.Contains(code, "func Test") {
+	// Don't modify complete functions or variable declarations
+	if strings.Contains(code, "func Test") ||
+		strings.Contains(code, "tests :") ||
+		strings.Contains(code, "var tests") {
 		return code
 	}
 
@@ -78,6 +80,17 @@ func parseGoTestCases(goCode string, maxCases int) ([]TestCase, error) {
 		return parseCompleteTestFunction(cleaned, maxCases)
 	}
 
+	// Some models output a declaration like:
+	//   tests := []struct{...}{...}
+	// Wrap it so it becomes valid Go and parse as a complete function.
+	if strings.Contains(cleaned, "tests :=") || strings.Contains(cleaned, "tests=") {
+		wrapped := fmt.Sprintf("package main\nfunc init() {\n%s\n}\n", cleaned)
+		cases, err := parseCompleteTestFunction(wrapped, maxCases)
+		if err == nil && len(cases) > 0 {
+			return cases, nil
+		}
+	}
+
 	// Fallback: parse as just test case array (old approach)
 	return parseTestCaseArray(cleaned, maxCases)
 }