use gofumpt

10 files changed, 876 insertions, 666 deletions

diff --git a/internal/llm/copilot.go b/internal/llm/copilot.go
index 16eeda6..d3b1a9d 100644
--- a/internal/llm/copilot.go
+++ b/internal/llm/copilot.go
@@ -4,6 +4,7 @@ package llm
 import (
 	"bytes"
 	"context"
+	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -13,7 +14,6 @@ import (
 	"strings"
 	"time"
 
-	"encoding/base64"
 	appver "codeberg.org/snonux/hexai/internal"
 	"codeberg.org/snonux/hexai/internal/logging"
 )
@@ -162,10 +162,14 @@ func buildCopilotChatRequest(o Options, messages []Message, defaultTemp *float64
 }
 
 func (c copilotClient) postJSON(ctx context.Context, url string, body []byte, headers map[string]string) (*http.Response, error) {
-    req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
-    if err != nil { return nil, err }
-    for k, v := range headers { req.Header.Set(k, v) }
-    return c.httpClient.Do(req)
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
+	if err != nil {
+		return nil, err
+	}
+	for k, v := range headers {
+		req.Header.Set(k, v)
+	}
+	return c.httpClient.Do(req)
 }
 
 func handleCopilotNon2xx(resp *http.Response, start time.Time) error {
@@ -194,55 +198,73 @@ func decodeCopilotChat(resp *http.Response, start time.Time) (copilotChatRespons
 // --- Copilot session token management ---
 
 type ghCopilotTokenResp struct {
-    Token string `json:"token"`
+	Token string `json:"token"`
 }
 
 func (c *copilotClient) ensureSession(ctx context.Context) error {
-    // If token valid for >60s, reuse
-    if c.sessionToken != "" && time.Now().Add(60*time.Second).Before(c.tokenExpiry) {
-        return nil
-    }
-    if strings.TrimSpace(c.apiKey) == "" {
-        return errors.New("missing Copilot API key")
-    }
-    req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://api.github.com/copilot_internal/v2/token", nil)
-    if err != nil { return err }
-    req.Header.Set("Authorization", "Bearer "+c.apiKey)
-    req.Header.Set("Accept", "application/json")
-    req.Header.Set("User-Agent", "hexai/"+appver.Version)
-    resp, err := c.httpClient.Do(req)
-    if err != nil { return err }
-    defer resp.Body.Close()
-    if resp.StatusCode < 200 || resp.StatusCode >= 300 {
-        return fmt.Errorf("copilot token http error: %d", resp.StatusCode)
-    }
-    var out ghCopilotTokenResp
-    if err := json.NewDecoder(resp.Body).Decode(&out); err != nil { return err }
-    if strings.TrimSpace(out.Token) == "" { return errors.New("empty copilot session token") }
-    // Parse JWT exp
-    exp := parseJWTExp(out.Token)
-    if exp.IsZero() { exp = time.Now().Add(10 * time.Minute) }
-    c.sessionToken = out.Token
-    c.tokenExpiry = exp
-    return nil
+	// If token valid for >60s, reuse
+	if c.sessionToken != "" && time.Now().Add(60*time.Second).Before(c.tokenExpiry) {
+		return nil
+	}
+	if strings.TrimSpace(c.apiKey) == "" {
+		return errors.New("missing Copilot API key")
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://api.github.com/copilot_internal/v2/token", nil)
+	if err != nil {
+		return err
+	}
+	req.Header.Set("Authorization", "Bearer "+c.apiKey)
+	req.Header.Set("Accept", "application/json")
+	req.Header.Set("User-Agent", "hexai/"+appver.Version)
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		return fmt.Errorf("copilot token http error: %d", resp.StatusCode)
+	}
+	var out ghCopilotTokenResp
+	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
+		return err
+	}
+	if strings.TrimSpace(out.Token) == "" {
+		return errors.New("empty copilot session token")
+	}
+	// Parse JWT exp
+	exp := parseJWTExp(out.Token)
+	if exp.IsZero() {
+		exp = time.Now().Add(10 * time.Minute)
+	}
+	c.sessionToken = out.Token
+	c.tokenExpiry = exp
+	return nil
 }
 
 var jwtExpRe = regexp.MustCompile(`"exp"\s*:\s*([0-9]+)`) // fallback if we can't base64 decode
 
 func parseJWTExp(token string) time.Time {
-    parts := strings.Split(token, ".")
-    if len(parts) < 2 { return time.Time{} }
-    b, err := base64.RawURLEncoding.DecodeString(parts[1])
-    if err != nil {
-        if m := jwtExpRe.FindStringSubmatch(token); len(m) == 2 {
-            if n, err2 := parseInt64(m[1]); err2 == nil { return time.Unix(n, 0) }
-        }
-        return time.Time{}
-    }
-    var payload struct{ Exp int64 `json:"exp"` }
-    _ = json.Unmarshal(b, &payload)
-    if payload.Exp == 0 { return time.Time{} }
-    return time.Unix(payload.Exp, 0)
+	parts := strings.Split(token, ".")
+	if len(parts) < 2 {
+		return time.Time{}
+	}
+	b, err := base64.RawURLEncoding.DecodeString(parts[1])
+	if err != nil {
+		if m := jwtExpRe.FindStringSubmatch(token); len(m) == 2 {
+			if n, err2 := parseInt64(m[1]); err2 == nil {
+				return time.Unix(n, 0)
+			}
+		}
+		return time.Time{}
+	}
+	var payload struct {
+		Exp int64 `json:"exp"`
+	}
+	_ = json.Unmarshal(b, &payload)
+	if payload.Exp == 0 {
+		return time.Time{}
+	}
+	return time.Unix(payload.Exp, 0)
 }
 
 func parseInt64(s string) (int64, error) { var n int64; _, err := fmt.Sscan(s, &n); return n, err }
@@ -250,99 +272,120 @@ func parseInt64(s string) (int64, error) { var n int64; _, err := fmt.Sscan(s, &
 // --- Copilot headers ---
 
 func (c *copilotClient) headersChat() map[string]string {
-    _ = c.ensureSession(context.Background())
-    h := map[string]string{
-        "Content-Type": "application/json; charset=utf-8",
-        "Accept":       "application/json",
-        "Authorization": "Bearer " + c.sessionToken,
-        "User-Agent":  "GitHubCopilotChat/0.8.0",
-        "Editor-Plugin-Version": "copilot-chat/0.8.0",
-        "Editor-Version":        "vscode/1.85.1",
-        "Openai-Intent":         "conversation-panel",
-        "Openai-Organization":   "github-copilot",
-        "VScode-MachineId":      randHex(64),
-        "VScode-SessionId":      randHex(8) + "-" + randHex(4) + "-" + randHex(4) + "-" + randHex(4) + "-" + randHex(12),
-        "X-Request-Id":          randHex(8) + "-" + randHex(4) + "-" + randHex(4) + "-" + randHex(4) + "-" + randHex(12),
-    }
-    return h
+	_ = c.ensureSession(context.Background())
+	h := map[string]string{
+		"Content-Type":          "application/json; charset=utf-8",
+		"Accept":                "application/json",
+		"Authorization":         "Bearer " + c.sessionToken,
+		"User-Agent":            "GitHubCopilotChat/0.8.0",
+		"Editor-Plugin-Version": "copilot-chat/0.8.0",
+		"Editor-Version":        "vscode/1.85.1",
+		"Openai-Intent":         "conversation-panel",
+		"Openai-Organization":   "github-copilot",
+		"VScode-MachineId":      randHex(64),
+		"VScode-SessionId":      randHex(8) + "-" + randHex(4) + "-" + randHex(4) + "-" + randHex(4) + "-" + randHex(12),
+		"X-Request-Id":          randHex(8) + "-" + randHex(4) + "-" + randHex(4) + "-" + randHex(4) + "-" + randHex(12),
+	}
+	return h
 }
 
 func (c *copilotClient) headersGhost() map[string]string {
-    _ = c.ensureSession(context.Background())
-    h := map[string]string{
-        "Content-Type": "application/json; charset=utf-8",
-        "Accept":       "*/*",
-        "Authorization": "Bearer " + c.sessionToken,
-        "User-Agent":  "GithubCopilot/1.155.0",
-        "Editor-Plugin-Version": "copilot/1.155.0",
-        "Editor-Version":        "vscode/1.85.1",
-        "Openai-Intent":         "copilot-ghost",
-        "Openai-Organization":   "github-copilot",
-        "VScode-MachineId":      randHex(64),
-        "VScode-SessionId":      randHex(8) + "-" + randHex(4) + "-" + randHex(4) + "-" + randHex(4) + "-" + randHex(12),
-        "X-Request-Id":          randHex(8) + "-" + randHex(4) + "-" + randHex(4) + "-" + randHex(4) + "-" + randHex(12),
-    }
-    return h
+	_ = c.ensureSession(context.Background())
+	h := map[string]string{
+		"Content-Type":          "application/json; charset=utf-8",
+		"Accept":                "*/*",
+		"Authorization":         "Bearer " + c.sessionToken,
+		"User-Agent":            "GithubCopilot/1.155.0",
+		"Editor-Plugin-Version": "copilot/1.155.0",
+		"Editor-Version":        "vscode/1.85.1",
+		"Openai-Intent":         "copilot-ghost",
+		"Openai-Organization":   "github-copilot",
+		"VScode-MachineId":      randHex(64),
+		"VScode-SessionId":      randHex(8) + "-" + randHex(4) + "-" + randHex(4) + "-" + randHex(4) + "-" + randHex(12),
+		"X-Request-Id":          randHex(8) + "-" + randHex(4) + "-" + randHex(4) + "-" + randHex(4) + "-" + randHex(12),
+	}
+	return h
 }
 
 func randHex(n int) string {
-    const hex = "0123456789abcdef"
-    b := make([]byte, n)
-    for i := range b {
-        b[i] = hex[int(time.Now().UnixNano()+int64(i))%len(hex)]
-    }
-    return string(b)
+	const hex = "0123456789abcdef"
+	b := make([]byte, n)
+	for i := range b {
+		b[i] = hex[int(time.Now().UnixNano()+int64(i))%len(hex)]
+	}
+	return string(b)
 }
 
 // --- Codex-style code completion ---
 
 // CodeCompletion implements CodeCompleter; returns up to n suggestions.
 func (c copilotClient) CodeCompletion(ctx context.Context, prompt string, suffix string, n int, language string, temperature float64) ([]string, error) {
-    if strings.TrimSpace(c.apiKey) == "" { return nil, errors.New("missing Copilot API key") }
-    if err := c.ensureSession(ctx); err != nil { return nil, err }
-    if n <= 0 { n = 1 }
-    maxTokens := 500
-    body := map[string]any{
-        "extra": map[string]any{
-            "language": language,
-            "next_indent": 0,
-            "prompt_tokens": 500,
-            "suffix_tokens": 400,
-            "trim_by_indentation": true,
-        },
-        "max_tokens": maxTokens,
-        "n": n,
-        "nwo": "hexai",
-        "prompt": prompt,
-        "stop": []string{"\n\n"},
-        "stream": true,
-        "suffix": suffix,
-        "temperature": temperature,
-        "top_p": 1,
-    }
-    buf, _ := json.Marshal(body)
-    url := "https://copilot-proxy.githubusercontent.com/v1/engines/copilot-codex/completions"
-    resp, err := c.postJSON(ctx, url, buf, c.headersGhost())
-    if err != nil { return nil, err }
-    defer resp.Body.Close()
-    if resp.StatusCode < 200 || resp.StatusCode >= 300 {
-        return nil, fmt.Errorf("copilot codex http error: %d", resp.StatusCode)
-    }
-    // Read all and parse lines that start with "data: " accumulating by index
-    raw, _ := io.ReadAll(resp.Body)
-    byIndex := make(map[int]string)
-    lines := strings.Split(string(raw), "\n")
-    for _, ln := range lines {
-        if !strings.HasPrefix(ln, "data: ") { continue }
-        var evt struct{ Choices []struct{ Index int `json:"index"`; Text string `json:"text"` } `json:"choices"` }
-        if err := json.Unmarshal([]byte(strings.TrimPrefix(ln, "data: ")), &evt); err != nil { continue }
-        for _, ch := range evt.Choices { byIndex[ch.Index] += ch.Text }
-    }
-    out := make([]string, 0, len(byIndex))
-    for i := 0; i < n; i++ {
-        if s, ok := byIndex[i]; ok && strings.TrimSpace(s) != "" { out = append(out, s) }
-    }
-    return out, nil
+	if strings.TrimSpace(c.apiKey) == "" {
+		return nil, errors.New("missing Copilot API key")
+	}
+	if err := c.ensureSession(ctx); err != nil {
+		return nil, err
+	}
+	if n <= 0 {
+		n = 1
+	}
+	maxTokens := 500
+	body := map[string]any{
+		"extra": map[string]any{
+			"language":            language,
+			"next_indent":         0,
+			"prompt_tokens":       500,
+			"suffix_tokens":       400,
+			"trim_by_indentation": true,
+		},
+		"max_tokens":  maxTokens,
+		"n":           n,
+		"nwo":         "hexai",
+		"prompt":      prompt,
+		"stop":        []string{"\n\n"},
+		"stream":      true,
+		"suffix":      suffix,
+		"temperature": temperature,
+		"top_p":       1,
+	}
+	buf, _ := json.Marshal(body)
+	url := "https://copilot-proxy.githubusercontent.com/v1/engines/copilot-codex/completions"
+	resp, err := c.postJSON(ctx, url, buf, c.headersGhost())
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		return nil, fmt.Errorf("copilot codex http error: %d", resp.StatusCode)
+	}
+	// Read all and parse lines that start with "data: " accumulating by index
+	raw, _ := io.ReadAll(resp.Body)
+	byIndex := make(map[int]string)
+	lines := strings.Split(string(raw), "\n")
+	for _, ln := range lines {
+		if !strings.HasPrefix(ln, "data: ") {
+			continue
+		}
+		var evt struct {
+			Choices []struct {
+				Index int    `json:"index"`
+				Text  string `json:"text"`
+			} `json:"choices"`
+		}
+		if err := json.Unmarshal([]byte(strings.TrimPrefix(ln, "data: ")), &evt); err != nil {
+			continue
+		}
+		for _, ch := range evt.Choices {
+			byIndex[ch.Index] += ch.Text
+		}
+	}
+	out := make([]string, 0, len(byIndex))
+	for i := 0; i < n; i++ {
+		if s, ok := byIndex[i]; ok && strings.TrimSpace(s) != "" {
+			out = append(out, s)
+		}
+	}
+	return out, nil
 }
 
 // newLineDataReader wraps a streaming body and exposes a JSON decoder that
diff --git a/internal/llm/copilot_http_test.go b/internal/llm/copilot_http_test.go
index 180e43e..d66311c 100644
--- a/internal/llm/copilot_http_test.go
+++ b/internal/llm/copilot_http_test.go
@@ -1,205 +1,261 @@
 package llm
 
 import (
-    "context"
-    "encoding/json"
-    "io"
-    "net/http"
-    "net/http/httptest"
-    "strings"
-    "testing"
-    "time"
-    "encoding/base64"
-    "os"
+	"context"
+	"encoding/base64"
+	"encoding/json"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"strings"
+	"testing"
+	"time"
 )
 
 type rtFunc2 func(*http.Request) (*http.Response, error)
+
 func (f rtFunc2) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
 
 func TestCopilot_EnsureSession_AndChat_Success(t *testing.T) {
-    if os.Getenv("HEXAI_TEST_SKIP_NET") == "1" { t.Skip("skip network-bound tests in restricted environments") }
-    // Mock chat endpoint
-    chatSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-        if r.URL.Path != "/chat/completions" { t.Fatalf("unexpected path: %s", r.URL.Path) }
-        _ = json.NewEncoder(w).Encode(map[string]any{"choices": []map[string]any{{"index":0, "message": map[string]string{"role":"assistant","content":"OK"}}}})
-    }))
-    defer chatSrv.Close()
-    c := newCopilot(chatSrv.URL, "gpt-4o-mini", "APIKEY", f64p(0.1)).(copilotClient)
-    // Intercept token endpoint to return a session token
-    tr := rtFunc2(func(r *http.Request) (*http.Response, error) {
-        if r.URL.Host == "api.github.com" && r.URL.Path == "/copilot_internal/v2/token" {
-            rw := httptest.NewRecorder()
-            _ = json.NewEncoder(rw).Encode(map[string]string{"token":"tok"})
-            res := rw.Result()
-            res.StatusCode = 200
-            return res, nil
-        }
-        // Fallback to default transport for chatSrv
-        return http.DefaultTransport.RoundTrip(r)
-    })
-    c.httpClient = &http.Client{Transport: tr, Timeout: 5 * time.Second}
-    out, err := c.Chat(context.Background(), []Message{{Role:"user", Content:"hi"}})
-    if err != nil || out != "OK" { t.Fatalf("copilot chat failed: %v %q", err, out) }
+	if os.Getenv("HEXAI_TEST_SKIP_NET") == "1" {
+		t.Skip("skip network-bound tests in restricted environments")
+	}
+	// Mock chat endpoint
+	chatSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != "/chat/completions" {
+			t.Fatalf("unexpected path: %s", r.URL.Path)
+		}
+		_ = json.NewEncoder(w).Encode(map[string]any{"choices": []map[string]any{{"index": 0, "message": map[string]string{"role": "assistant", "content": "OK"}}}})
+	}))
+	defer chatSrv.Close()
+	c := newCopilot(chatSrv.URL, "gpt-4o-mini", "APIKEY", f64p(0.1)).(copilotClient)
+	// Intercept token endpoint to return a session token
+	tr := rtFunc2(func(r *http.Request) (*http.Response, error) {
+		if r.URL.Host == "api.github.com" && r.URL.Path == "/copilot_internal/v2/token" {
+			rw := httptest.NewRecorder()
+			_ = json.NewEncoder(rw).Encode(map[string]string{"token": "tok"})
+			res := rw.Result()
+			res.StatusCode = 200
+			return res, nil
+		}
+		// Fallback to default transport for chatSrv
+		return http.DefaultTransport.RoundTrip(r)
+	})
+	c.httpClient = &http.Client{Transport: tr, Timeout: 5 * time.Second}
+	out, err := c.Chat(context.Background(), []Message{{Role: "user", Content: "hi"}})
+	if err != nil || out != "OK" {
+		t.Fatalf("copilot chat failed: %v %q", err, out)
+	}
 }
 
 func TestCopilot_HandleNon2xx(t *testing.T) {
-    b, _ := json.Marshal(map[string]any{"error": map[string]any{"message":"bad","type":"invalid"}})
-    resp := &http.Response{StatusCode: 400, Body: io.NopCloser(bytesReader(b))}
-    if err := handleCopilotNon2xx(resp, time.Now()); err == nil { t.Fatalf("expected error") }
+	b, _ := json.Marshal(map[string]any{"error": map[string]any{"message": "bad", "type": "invalid"}})
+	resp := &http.Response{StatusCode: 400, Body: io.NopCloser(bytesReader(b))}
+	if err := handleCopilotNon2xx(resp, time.Now()); err == nil {
+		t.Fatalf("expected error")
+	}
 }
 
 func TestCopilot_CodeCompletion_Success(t *testing.T) {
-    c := newCopilot("https://api.githubcopilot.com", "gpt-4o-mini", "API", f64p(0.1)).(copilotClient)
-    tr := rtFunc2(func(r *http.Request) (*http.Response, error) {
-        // Token endpoint
-        if r.URL.Host == "api.github.com" && r.URL.Path == "/copilot_internal/v2/token" {
-            rw := httptest.NewRecorder()
-            _ = json.NewEncoder(rw).Encode(map[string]string{"token":"tok"})
-            res := rw.Result(); res.StatusCode = 200; return res, nil
-        }
-        // Codex completion endpoint
-        if r.URL.Host == "copilot-proxy.githubusercontent.com" && strings.HasSuffix(r.URL.Path, "/v1/engines/copilot-codex/completions") {
-            rw := httptest.NewRecorder()
-            // two choices for index 0 and 1
-            rw.WriteString("data: {\"choices\":[{\"index\":0,\"text\":\"A\"}]}\n")
-            rw.WriteString("data: {\"choices\":[{\"index\":1,\"text\":\"B\"}]}\n")
-            res := rw.Result(); res.StatusCode = 200; return res, nil
-        }
-        return http.DefaultTransport.RoundTrip(r)
-    })
-    c.httpClient = &http.Client{Transport: tr, Timeout: 5 * time.Second}
-    out, err := c.CodeCompletion(context.Background(), "p", "s", 2, "go", 0.1)
-    if err != nil || len(out) != 2 || out[0] != "A" || out[1] != "B" {
-        t.Fatalf("codex: %v %#v", err, out)
-    }
+	c := newCopilot("https://api.githubcopilot.com", "gpt-4o-mini", "API", f64p(0.1)).(copilotClient)
+	tr := rtFunc2(func(r *http.Request) (*http.Response, error) {
+		// Token endpoint
+		if r.URL.Host == "api.github.com" && r.URL.Path == "/copilot_internal/v2/token" {
+			rw := httptest.NewRecorder()
+			_ = json.NewEncoder(rw).Encode(map[string]string{"token": "tok"})
+			res := rw.Result()
+			res.StatusCode = 200
+			return res, nil
+		}
+		// Codex completion endpoint
+		if r.URL.Host == "copilot-proxy.githubusercontent.com" && strings.HasSuffix(r.URL.Path, "/v1/engines/copilot-codex/completions") {
+			rw := httptest.NewRecorder()
+			// two choices for index 0 and 1
+			rw.WriteString("data: {\"choices\":[{\"index\":0,\"text\":\"A\"}]}\n")
+			rw.WriteString("data: {\"choices\":[{\"index\":1,\"text\":\"B\"}]}\n")
+			res := rw.Result()
+			res.StatusCode = 200
+			return res, nil
+		}
+		return http.DefaultTransport.RoundTrip(r)
+	})
+	c.httpClient = &http.Client{Transport: tr, Timeout: 5 * time.Second}
+	out, err := c.CodeCompletion(context.Background(), "p", "s", 2, "go", 0.1)
+	if err != nil || len(out) != 2 || out[0] != "A" || out[1] != "B" {
+		t.Fatalf("codex: %v %#v", err, out)
+	}
 }
 
 func TestCopilot_Chat_MultiChoice_And_ErrorBody(t *testing.T) {
-    if os.Getenv("HEXAI_TEST_SKIP_NET") == "1" { t.Skip("skip network-bound tests in restricted environments") }
-    // Chat multi-choice: return two choices; client returns first content
-    srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-        _ = json.NewEncoder(w).Encode(map[string]any{
-            "choices": []map[string]any{
-                {"index": 0, "finish_reason": "stop", "message": map[string]string{"role": "assistant", "content": "FIRST"}},
-                {"index": 1, "finish_reason": "length", "message": map[string]string{"role": "assistant", "content": "SECOND"}},
-            },
-        })
-    }))
-    defer srv.Close()
-    c := newCopilot(srv.URL, "gpt-4o-mini", "KEY", f64p(0.1)).(copilotClient)
-    // Token success
-    tr := rtFunc2(func(r *http.Request) (*http.Response, error) {
-        if r.URL.Host == "api.github.com" && r.URL.Path == "/copilot_internal/v2/token" {
-            rw := httptest.NewRecorder(); _ = json.NewEncoder(rw).Encode(map[string]string{"token":"tok"}); res := rw.Result(); res.StatusCode = 200; return res, nil
-        }
-        return http.DefaultTransport.RoundTrip(r)
-    })
-    c.httpClient = &http.Client{Transport: tr, Timeout: 5 * time.Second}
-    out, err := c.Chat(context.Background(), []Message{{Role: "user", Content: "hi"}})
-    if err != nil || out != "FIRST" { t.Fatalf("copilot multi-choice: %v %q", err, out) }
+	if os.Getenv("HEXAI_TEST_SKIP_NET") == "1" {
+		t.Skip("skip network-bound tests in restricted environments")
+	}
+	// Chat multi-choice: return two choices; client returns first content
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"choices": []map[string]any{
+				{"index": 0, "finish_reason": "stop", "message": map[string]string{"role": "assistant", "content": "FIRST"}},
+				{"index": 1, "finish_reason": "length", "message": map[string]string{"role": "assistant", "content": "SECOND"}},
+			},
+		})
+	}))
+	defer srv.Close()
+	c := newCopilot(srv.URL, "gpt-4o-mini", "KEY", f64p(0.1)).(copilotClient)
+	// Token success
+	tr := rtFunc2(func(r *http.Request) (*http.Response, error) {
+		if r.URL.Host == "api.github.com" && r.URL.Path == "/copilot_internal/v2/token" {
+			rw := httptest.NewRecorder()
+			_ = json.NewEncoder(rw).Encode(map[string]string{"token": "tok"})
+			res := rw.Result()
+			res.StatusCode = 200
+			return res, nil
+		}
+		return http.DefaultTransport.RoundTrip(r)
+	})
+	c.httpClient = &http.Client{Transport: tr, Timeout: 5 * time.Second}
+	out, err := c.Chat(context.Background(), []Message{{Role: "user", Content: "hi"}})
+	if err != nil || out != "FIRST" {
+		t.Fatalf("copilot multi-choice: %v %q", err, out)
+	}
 
-    // Non-2xx with error body
-    srv2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-        w.WriteHeader(403)
-        _ = json.NewEncoder(w).Encode(map[string]any{"error": map[string]any{"message":"denied","type":"forbidden"}})
-    }))
-    defer srv2.Close()
-    c2 := newCopilot(srv2.URL, "gpt-4o-mini", "KEY", f64p(0.1)).(copilotClient)
-    c2.httpClient = &http.Client{Transport: tr, Timeout: 5 * time.Second}
-    if _, err := c2.Chat(context.Background(), []Message{{Role:"user", Content:"hi"}}); err == nil {
-        t.Fatalf("expected error for copilot non-2xx with error body")
-    }
+	// Non-2xx with error body
+	srv2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(403)
+		_ = json.NewEncoder(w).Encode(map[string]any{"error": map[string]any{"message": "denied", "type": "forbidden"}})
+	}))
+	defer srv2.Close()
+	c2 := newCopilot(srv2.URL, "gpt-4o-mini", "KEY", f64p(0.1)).(copilotClient)
+	c2.httpClient = &http.Client{Transport: tr, Timeout: 5 * time.Second}
+	if _, err := c2.Chat(context.Background(), []Message{{Role: "user", Content: "hi"}}); err == nil {
+		t.Fatalf("expected error for copilot non-2xx with error body")
+	}
 }
 
 func TestCopilot_Chat_NoChoices_Error(t *testing.T) {
-    if os.Getenv("HEXAI_TEST_SKIP_NET") == "1" { t.Skip("skip network-bound tests in restricted environments") }
-    srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-        _ = json.NewEncoder(w).Encode(map[string]any{"choices": []any{}})
-    }))
-    defer srv.Close()
-    c := newCopilot(srv.URL, "gpt-4o-mini", "KEY", f64p(0.1)).(copilotClient)
-    tr := rtFunc2(func(r *http.Request) (*http.Response, error) {
-        if r.URL.Host == "api.github.com" && r.URL.Path == "/copilot_internal/v2/token" {
-            rw := httptest.NewRecorder(); _ = json.NewEncoder(rw).Encode(map[string]string{"token":"tok"}); res := rw.Result(); res.StatusCode = 200; return res, nil
-        }
-        return http.DefaultTransport.RoundTrip(r)
-    })
-    c.httpClient = &http.Client{Transport: tr, Timeout: 5 * time.Second}
-    if _, err := c.Chat(context.Background(), []Message{{Role:"user", Content:"hi"}}); err == nil {
-        t.Fatalf("expected error when no choices returned")
-    }
+	if os.Getenv("HEXAI_TEST_SKIP_NET") == "1" {
+		t.Skip("skip network-bound tests in restricted environments")
+	}
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_ = json.NewEncoder(w).Encode(map[string]any{"choices": []any{}})
+	}))
+	defer srv.Close()
+	c := newCopilot(srv.URL, "gpt-4o-mini", "KEY", f64p(0.1)).(copilotClient)
+	tr := rtFunc2(func(r *http.Request) (*http.Response, error) {
+		if r.URL.Host == "api.github.com" && r.URL.Path == "/copilot_internal/v2/token" {
+			rw := httptest.NewRecorder()
+			_ = json.NewEncoder(rw).Encode(map[string]string{"token": "tok"})
+			res := rw.Result()
+			res.StatusCode = 200
+			return res, nil
+		}
+		return http.DefaultTransport.RoundTrip(r)
+	})
+	c.httpClient = &http.Client{Transport: tr, Timeout: 5 * time.Second}
+	if _, err := c.Chat(context.Background(), []Message{{Role: "user", Content: "hi"}}); err == nil {
+		t.Fatalf("expected error when no choices returned")
+	}
 }
 
 func TestCopilot_Chat_DecodeError_StatusOK(t *testing.T) {
-    if os.Getenv("HEXAI_TEST_SKIP_NET") == "1" { t.Skip("skip network-bound tests in restricted environments") }
-    // Chat returns 200 but invalid JSON; expect decode error
-    srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-        io.WriteString(w, "{invalid")
-    }))
-    defer srv.Close()
-    c := newCopilot(srv.URL, "gpt-4o-mini", "KEY", f64p(0.1)).(copilotClient)
-    tr := rtFunc2(func(r *http.Request) (*http.Response, error) {
-        if r.URL.Host == "api.github.com" && r.URL.Path == "/copilot_internal/v2/token" {
-            rw := httptest.NewRecorder(); _ = json.NewEncoder(rw).Encode(map[string]string{"token":"tok"}); res := rw.Result(); res.StatusCode = 200; return res, nil
-        }
-        return http.DefaultTransport.RoundTrip(r)
-    })
-    c.httpClient = &http.Client{Transport: tr, Timeout: 5 * time.Second}
-    if _, err := c.Chat(context.Background(), []Message{{Role:"user", Content:"hi"}}); err == nil {
-        t.Fatalf("expected decode error for invalid body")
-    }
+	if os.Getenv("HEXAI_TEST_SKIP_NET") == "1" {
+		t.Skip("skip network-bound tests in restricted environments")
+	}
+	// Chat returns 200 but invalid JSON; expect decode error
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		io.WriteString(w, "{invalid")
+	}))
+	defer srv.Close()
+	c := newCopilot(srv.URL, "gpt-4o-mini", "KEY", f64p(0.1)).(copilotClient)
+	tr := rtFunc2(func(r *http.Request) (*http.Response, error) {
+		if r.URL.Host == "api.github.com" && r.URL.Path == "/copilot_internal/v2/token" {
+			rw := httptest.NewRecorder()
+			_ = json.NewEncoder(rw).Encode(map[string]string{"token": "tok"})
+			res := rw.Result()
+			res.StatusCode = 200
+			return res, nil
+		}
+		return http.DefaultTransport.RoundTrip(r)
+	})
+	c.httpClient = &http.Client{Transport: tr, Timeout: 5 * time.Second}
+	if _, err := c.Chat(context.Background(), []Message{{Role: "user", Content: "hi"}}); err == nil {
+		t.Fatalf("expected decode error for invalid body")
+	}
 }
 
 func TestCopilot_CodeCompletion_MalformedAndEmpty(t *testing.T) {
-    c := newCopilot("https://api.githubcopilot.com", "gpt-4o-mini", "API", f64p(0.1)).(copilotClient)
-    tr := rtFunc2(func(r *http.Request) (*http.Response, error) {
-        if r.URL.Host == "api.github.com" && r.URL.Path == "/copilot_internal/v2/token" {
-            rw := httptest.NewRecorder(); _ = json.NewEncoder(rw).Encode(map[string]string{"token":"tok"}); res := rw.Result(); res.StatusCode = 200; return res, nil
-        }
-        if r.URL.Host == "copilot-proxy.githubusercontent.com" && strings.HasSuffix(r.URL.Path, "/v1/engines/copilot-codex/completions") {
-            rw := httptest.NewRecorder()
-            // malformed line
-            rw.WriteString("data: {bad}\n")
-            // done; should produce empty suggestions
-            rw.WriteString("data: [DONE]\n")
-            res := rw.Result(); res.StatusCode = 200; return res, nil
-        }
-        return http.DefaultTransport.RoundTrip(r)
-    })
-    c.httpClient = &http.Client{Transport: tr, Timeout: 5 * time.Second}
-    out, err := c.CodeCompletion(context.Background(), "p", "s", 1, "go", 0.1)
-    if err != nil { t.Fatalf("unexpected error: %v", err) }
-    if len(out) != 0 { t.Fatalf("expected empty suggestions, got %#v", out) }
+	c := newCopilot("https://api.githubcopilot.com", "gpt-4o-mini", "API", f64p(0.1)).(copilotClient)
+	tr := rtFunc2(func(r *http.Request) (*http.Response, error) {
+		if r.URL.Host == "api.github.com" && r.URL.Path == "/copilot_internal/v2/token" {
+			rw := httptest.NewRecorder()
+			_ = json.NewEncoder(rw).Encode(map[string]string{"token": "tok"})
+			res := rw.Result()
+			res.StatusCode = 200
+			return res, nil
+		}
+		if r.URL.Host == "copilot-proxy.githubusercontent.com" && strings.HasSuffix(r.URL.Path, "/v1/engines/copilot-codex/completions") {
+			rw := httptest.NewRecorder()
+			// malformed line
+			rw.WriteString("data: {bad}\n")
+			// done; should produce empty suggestions
+			rw.WriteString("data: [DONE]\n")
+			res := rw.Result()
+			res.StatusCode = 200
+			return res, nil
+		}
+		return http.DefaultTransport.RoundTrip(r)
+	})
+	c.httpClient = &http.Client{Transport: tr, Timeout: 5 * time.Second}
+	out, err := c.CodeCompletion(context.Background(), "p", "s", 1, "go", 0.1)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if len(out) != 0 {
+		t.Fatalf("expected empty suggestions, got %#v", out)
+	}
 
-    // Now include one good chunk after malformed
-    tr2 := rtFunc2(func(r *http.Request) (*http.Response, error) {
-        if r.URL.Host == "api.github.com" && r.URL.Path == "/copilot_internal/v2/token" {
-            rw := httptest.NewRecorder(); _ = json.NewEncoder(rw).Encode(map[string]string{"token":"tok"}); res := rw.Result(); res.StatusCode = 200; return res, nil
-        }
-        if r.URL.Host == "copilot-proxy.githubusercontent.com" && strings.HasSuffix(r.URL.Path, "/v1/engines/copilot-codex/completions") {
-            rw := httptest.NewRecorder()
-            rw.WriteString("data: {bad}\n")
-            rw.WriteString("data: {\"choices\":[{\"index\":0,\"text\":\"OK\"}]}\n")
-            rw.WriteString("data: [DONE]\n")
-            res := rw.Result(); res.StatusCode = 200; return res, nil
-        }
-        return http.DefaultTransport.RoundTrip(r)
-    })
-    c.httpClient = &http.Client{Transport: tr2, Timeout: 5 * time.Second}
-    out2, err := c.CodeCompletion(context.Background(), "p", "s", 1, "go", 0.1)
-    if err != nil || len(out2) != 1 || out2[0] != "OK" { t.Fatalf("unexpected: %v %#v", err, out2) }
+	// Now include one good chunk after malformed
+	tr2 := rtFunc2(func(r *http.Request) (*http.Response, error) {
+		if r.URL.Host == "api.github.com" && r.URL.Path == "/copilot_internal/v2/token" {
+			rw := httptest.NewRecorder()
+			_ = json.NewEncoder(rw).Encode(map[string]string{"token": "tok"})
+			res := rw.Result()
+			res.StatusCode = 200
+			return res, nil
+		}
+		if r.URL.Host == "copilot-proxy.githubusercontent.com" && strings.HasSuffix(r.URL.Path, "/v1/engines/copilot-codex/completions") {
+			rw := httptest.NewRecorder()
+			rw.WriteString("data: {bad}\n")
+			rw.WriteString("data: {\"choices\":[{\"index\":0,\"text\":\"OK\"}]}\n")
+			rw.WriteString("data: [DONE]\n")
+			res := rw.Result()
+			res.StatusCode = 200
+			return res, nil
+		}
+		return http.DefaultTransport.RoundTrip(r)
+	})
+	c.httpClient = &http.Client{Transport: tr2, Timeout: 5 * time.Second}
+	out2, err := c.CodeCompletion(context.Background(), "p", "s", 1, "go", 0.1)
+	if err != nil || len(out2) != 1 || out2[0] != "OK" {
+		t.Fatalf("unexpected: %v %#v", err, out2)
+	}
 }
 
 func TestParseJWTExp_AndParseInt64(t *testing.T) {
-    // Valid base64 payload
-    payload := `{"exp": 1700000000}`
-    b := base64.RawURLEncoding.EncodeToString([]byte(payload))
-    tok := "x." + b + ".y"
-    if tm := parseJWTExp(tok); tm.IsZero() { t.Fatalf("expected non-zero time") }
-    if n, err := parseInt64("123"); err != nil || n != 123 { t.Fatalf("parseInt64: %v %d", err, n) }
+	// Valid base64 payload
+	payload := `{"exp": 1700000000}`
+	b := base64.RawURLEncoding.EncodeToString([]byte(payload))
+	tok := "x." + b + ".y"
+	if tm := parseJWTExp(tok); tm.IsZero() {
+		t.Fatalf("expected non-zero time")
+	}
+	if n, err := parseInt64("123"); err != nil || n != 123 {
+		t.Fatalf("parseInt64: %v %d", err, n)
+	}
 }
 
 // bytesReader wraps a byte slice with an io.ReadCloser without importing extra.
 type bytesReader []byte
+
 func (b bytesReader) Read(p []byte) (int, error) { n := copy(p, b); return n, io.EOF }
-func (b bytesReader) Close() error { return nil }
+func (b bytesReader) Close() error               { return nil }
diff --git a/internal/llm/ollama_test.go b/internal/llm/ollama_test.go
index 15f9cff..8bd33ca 100644
--- a/internal/llm/ollama_test.go
+++ b/internal/llm/ollama_test.go
@@ -1,173 +1,217 @@
 package llm
 
 import (
-    "context"
-    "encoding/json"
-    "fmt"
-    "net/http"
-    "net/http/httptest"
-    "strings"
-    "testing"
-    "time"
-    "os"
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"strings"
+	"testing"
+	"time"
 )
 
 func TestBuildOllamaRequest_OptionsAndStream(t *testing.T) {
-    o := Options{Model: "codemodel", Temperature: 0, MaxTokens: 256, Stop: []string{"STOP"}}
-    msgs := []Message{{Role: "user", Content: "hello"}}
-    req := buildOllamaRequest(o, msgs, f64p(0.2), false)
-    if req.Model != "codemodel" || req.Stream { t.Fatalf("model/stream mismatch: %+v", req) }
-    if req.Options == nil { t.Fatalf("expected options map") }
-    if req.Options.(map[string]any)["temperature"].(float64) != 0.2 { t.Fatalf("default temp not applied") }
-    if req.Options.(map[string]any)["num_predict"].(int) != 256 { t.Fatalf("num_predict not applied") }
-    if req.Options.(map[string]any)["stop"].([]string)[0] != "STOP" { t.Fatalf("stop not applied") }
-
-    req2 := buildOllamaRequest(o, msgs, f64p(0.2), true)
-    if !req2.Stream { t.Fatalf("expected stream=true") }
+	o := Options{Model: "codemodel", Temperature: 0, MaxTokens: 256, Stop: []string{"STOP"}}
+	msgs := []Message{{Role: "user", Content: "hello"}}
+	req := buildOllamaRequest(o, msgs, f64p(0.2), false)
+	if req.Model != "codemodel" || req.Stream {
+		t.Fatalf("model/stream mismatch: %+v", req)
+	}
+	if req.Options == nil {
+		t.Fatalf("expected options map")
+	}
+	if req.Options.(map[string]any)["temperature"].(float64) != 0.2 {
+		t.Fatalf("default temp not applied")
+	}
+	if req.Options.(map[string]any)["num_predict"].(int) != 256 {
+		t.Fatalf("num_predict not applied")
+	}
+	if req.Options.(map[string]any)["stop"].([]string)[0] !=