jellybit/internal/llm/llm_test.go

package llm

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
)

// newTestProvider собирает openai-compat клиент на адрес стенда и убирает
// паузы между ретраями, чтобы тесты не висели.
func newTestProvider(t *testing.T, baseURL, apiKey string) *openAICompat {
	t.Helper()
	p, err := newOpenAICompat(Config{
		Type:    "openai-compat",
		BaseURL: baseURL,
		APIKey:  apiKey,
		Model:   "test-model",
	}, nil)
	if err != nil {
		t.Fatalf("newOpenAICompat: %v", err)
	}
	p.retryWait = 0
	return p
}

func TestComplete_Success(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if got := r.Header.Get("Authorization"); got != "Bearer secret" {
			t.Errorf("Authorization = %q, want Bearer secret", got)
		}
		body, _ := io.ReadAll(r.Body)
		var req chatRequest
		if err := json.Unmarshal(body, &req); err != nil {
			t.Fatalf("decode request: %v", err)
		}
		if req.Model != "test-model" {
			t.Errorf("model = %q, want test-model", req.Model)
		}
		if req.ResponseFormat == nil || req.ResponseFormat.Type != "json_object" {
			t.Errorf("response_format = %+v, want json_object", req.ResponseFormat)
		}
		_, _ = io.WriteString(w, `{"model":"resolved-model",
			"choices":[{"message":{"content":"{\"ok\":true}"},"finish_reason":"stop"}],
			"usage":{"prompt_tokens":10,"completion_tokens":3,"total_tokens":13,"cost":0.0001}}`)
	}))
	defer srv.Close()

	p := newTestProvider(t, srv.URL, "secret")
	resp, err := p.Complete(context.Background(), Request{
		Messages: []Message{{Role: RoleUser, Content: "hi"}},
		JSONMode: true,
	})
	if err != nil {
		t.Fatalf("Complete: %v", err)
	}
	if resp.Content != `{"ok":true}` {
		t.Errorf("content = %q", resp.Content)
	}
	if resp.Model != "resolved-model" {
		t.Errorf("model = %q", resp.Model)
	}
	if resp.Usage.TotalTokens != 13 || resp.Usage.Cost != 0.0001 {
		t.Errorf("usage = %+v", resp.Usage)
	}
}

func TestComplete_NoJSONModeOmitsResponseFormat(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		var req chatRequest
		_ = json.Unmarshal(body, &req)
		if req.ResponseFormat != nil {
			t.Errorf("response_format should be omitted, got %+v", req.ResponseFormat)
		}
		if !strings.Contains(string(body), `"temperature":0`) {
			t.Errorf("temperature 0 must be sent explicitly, body: %s", body)
		}
		_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"hello"}}]}`)
	}))
	defer srv.Close()

	p := newTestProvider(t, srv.URL, "")
	zero := 0.0
	resp, err := p.Complete(context.Background(), Request{
		Messages:    []Message{{Role: RoleUser, Content: "hi"}},
		Temperature: &zero,
	})
	if err != nil {
		t.Fatalf("Complete: %v", err)
	}
	if resp.Content != "hello" {
		t.Errorf("content = %q", resp.Content)
	}
}

func TestComplete_RetriesOn5xxThenSucceeds(t *testing.T) {
	var calls atomic.Int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if calls.Add(1) <= 2 {
			w.WriteHeader(http.StatusBadGateway)
			_, _ = io.WriteString(w, "upstream down")
			return
		}
		_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"recovered"}}]}`)
	}))
	defer srv.Close()

	p := newTestProvider(t, srv.URL, "")
	resp, err := p.Complete(context.Background(), Request{
		Messages: []Message{{Role: RoleUser, Content: "hi"}},
	})
	if err != nil {
		t.Fatalf("Complete: %v", err)
	}
	if resp.Content != "recovered" {
		t.Errorf("content = %q", resp.Content)
	}
	if got := calls.Load(); got != 3 {
		t.Errorf("calls = %d, want 3", got)
	}
}

func TestComplete_429IsRetryable(t *testing.T) {
	var calls atomic.Int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if calls.Add(1) == 1 {
			w.WriteHeader(http.StatusTooManyRequests)
			return
		}
		_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"ok"}}]}`)
	}))
	defer srv.Close()

	p := newTestProvider(t, srv.URL, "")
	if _, err := p.Complete(context.Background(), Request{
		Messages: []Message{{Role: RoleUser, Content: "hi"}},
	}); err != nil {
		t.Fatalf("Complete: %v", err)
	}
	if got := calls.Load(); got != 2 {
		t.Errorf("calls = %d, want 2", got)
	}
}

func TestComplete_4xxNotRetried(t *testing.T) {
	var calls atomic.Int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls.Add(1)
		w.WriteHeader(http.StatusBadRequest)
		_, _ = io.WriteString(w, `{"error":{"message":"bad model"}}`)
	}))
	defer srv.Close()

	p := newTestProvider(t, srv.URL, "")
	_, err := p.Complete(context.Background(), Request{
		Messages: []Message{{Role: RoleUser, Content: "hi"}},
	})
	if err == nil {
		t.Fatal("Complete: want error on 400")
	}
	if !strings.Contains(err.Error(), "status 400") {
		t.Errorf("err = %v, want status 400", err)
	}
	if got := calls.Load(); got != 1 {
		t.Errorf("calls = %d, want 1 (no retry on 4xx)", got)
	}
}

func TestComplete_ExhaustsRetries(t *testing.T) {
	var calls atomic.Int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls.Add(1)
		w.WriteHeader(http.StatusServiceUnavailable)
	}))
	defer srv.Close()

	p := newTestProvider(t, srv.URL, "")
	_, err := p.Complete(context.Background(), Request{
		Messages: []Message{{Role: RoleUser, Content: "hi"}},
	})
	if err == nil {
		t.Fatal("want error after exhausting retries")
	}
	if got := calls.Load(); got != maxAttempts {
		t.Errorf("calls = %d, want %d", got, maxAttempts)
	}
}

func TestComplete_ProviderErrorInBody(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// HTTP 200, но ошибка в теле — частый паттерн шлюзов.
		_, _ = io.WriteString(w, `{"error":{"message":"context length exceeded"}}`)
	}))
	defer srv.Close()

	p := newTestProvider(t, srv.URL, "")
	_, err := p.Complete(context.Background(), Request{
		Messages: []Message{{Role: RoleUser, Content: "hi"}},
	})
	if err == nil || !strings.Contains(err.Error(), "context length exceeded") {
		t.Fatalf("err = %v, want provider error", err)
	}
}

func TestComplete_EmptyMessages(t *testing.T) {
	p := newTestProvider(t, "http://example.invalid", "")
	if _, err := p.Complete(context.Background(), Request{}); err == nil {
		t.Fatal("want error on empty messages")
	}
}

func TestNew_UnknownType(t *testing.T) {
	if _, err := New(Config{Type: "anthropic", Model: "x", BaseURL: "http://x"}, nil); err == nil {
		t.Fatal("want error for unknown type")
	}
	if _, err := New(Config{Type: ""}, nil); err == nil {
		t.Fatal("want error for empty type")
	}
}

func TestNew_OpenAICompatValidation(t *testing.T) {
	if _, err := New(Config{Type: "openai-compat", Model: "x"}, nil); err == nil {
		t.Fatal("want error for empty base_url")
	}
	if _, err := New(Config{Type: "openai-compat", BaseURL: "http://x"}, nil); err == nil {
		t.Fatal("want error for empty model")
	}
	if _, err := New(Config{Type: "openai-compat", BaseURL: "http://x", Model: "m"}, nil); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}

// Признак обрыва генерации по длине доходит до вызывающего: по одному телу
// ответа обрыв неотличим от мусора, а повторять такой запрос бессмысленно.
func TestComplete_FinishReasonReachesCaller(t *testing.T) {
	for _, want := range []string{"length", "stop"} {
		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
			_, _ = io.WriteString(w, `{"model":"m","choices":[{"message":{"content":"{\"type\":"},`+
				`"finish_reason":"`+want+`"}]}`)
		}))

		p := newTestProvider(t, srv.URL, "")
		resp, err := p.Complete(context.Background(), Request{
			Messages: []Message{{Role: RoleUser, Content: "hi"}},
		})
		srv.Close()
		if err != nil {
			t.Fatalf("Complete: %v", err)
		}
		if resp.FinishReason != want {
			t.Errorf("finish_reason = %q, want %q", resp.FinishReason, want)
		}
	}
	if FinishLength != "length" {
		t.Errorf("FinishLength = %q, want length", FinishLength)
	}
}

// Отказ по размеру запроса помечается sentinel'ом там, где ещё целы
// HTTP-статус и ПОЛНОЕ тело: вызывающему нечего матчить по тексту, а текст
// итоговой ошибки обрезан (snippet) и несёт эхо запроса.
func TestComplete_RequestTooLargeSentinel(t *testing.T) {
	cases := []struct {
		name   string
		status int
		body   string
		want   bool
	}{
		{"413 без разбора текста", http.StatusRequestEntityTooLarge, `payload too big`, true},
		{"400 с маркером в теле", http.StatusBadRequest,
			`{"error":{"message":"This model's maximum context length is 128000 tokens"}}`, true},
		{"400 с маркером за пределом snippet", http.StatusBadRequest,
			`{"echo":"` + strings.Repeat("x", 400) + `","error":{"message":"prompt is too long"}}`, true},
		{"400 по другой причине", http.StatusBadRequest,
			`{"error":{"message":"invalid api key"}}`, false},
		{"500 — транзиентный сбой, не размер", http.StatusInternalServerError,
			`context length exceeded`, false},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
				w.WriteHeader(tc.status)
				_, _ = io.WriteString(w, tc.body)
			}))
			defer srv.Close()

			p := newTestProvider(t, srv.URL, "")
			_, err := p.Complete(context.Background(), Request{
				Messages: []Message{{Role: RoleUser, Content: "hi"}},
			})
			if err == nil {
				t.Fatal("ожидалась ошибка")
			}
			if got := errors.Is(err, ErrRequestTooLarge); got != tc.want {
				t.Errorf("errors.Is(err, ErrRequestTooLarge) = %v, want %v (err = %v)",
					got, tc.want, err)
			}
		})
	}
}