zzet--gortex/internal/embedding/api_test.go

package embedding

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestAPIProvider_Concurrent asserts the API provider declares itself
// safe to call concurrently — the signal the embedding pool gates on.
func TestAPIProvider_Concurrent(t *testing.T) {
	p := NewAPIProvider("http://localhost:11434", "")
	assert.True(t, p.Concurrent(), "the API provider must opt into concurrent embedding")
}

// TestParseRetryAfter covers the delta-seconds Retry-After parser.
func TestParseRetryAfter(t *testing.T) {
	assert.Equal(t, 12*time.Second, parseRetryAfter("12"))
	assert.Equal(t, time.Duration(0), parseRetryAfter(""))
	assert.Equal(t, time.Duration(0), parseRetryAfter("  "))
	assert.Equal(t, time.Duration(0), parseRetryAfter("Wed, 21 Oct 2026 07:28:00 GMT"),
		"the HTTP-date form is not parsed — returns 0 so the caller uses a fixed backoff")
	assert.Equal(t, time.Duration(0), parseRetryAfter("-5"), "a negative hint is rejected")
}

// TestAPIProvider_HonorsRetryAfterOn429 asserts the provider retries
// once after an HTTP 429, honouring the Retry-After header, and
// succeeds when the retry returns 200.
func TestAPIProvider_HonorsRetryAfterOn429(t *testing.T) {
	var calls int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		n := atomic.AddInt32(&calls, 1)
		if n == 1 {
			// First call: rate-limited with a 1-second hint.
			w.Header().Set("Retry-After", "1")
			w.WriteHeader(http.StatusTooManyRequests)
			return
		}
		// Retry: succeed with an OpenAI-shaped embedding response.
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(openAIResponse{
			Data: []openAIEmbedding{{Embedding: []float32{1, 2, 3}, Index: 0}},
		})
	}))
	defer srv.Close()

	// srv.URL has no Ollama markers, so the provider uses OpenAI format.
	p := NewAPIProvider(srv.URL, "text-embedding-3-small")

	start := time.Now()
	vecs, err := p.EmbedBatch(context.Background(), []string{"hello"})
	require.NoError(t, err, "the embed must succeed after the post-429 retry")
	require.Len(t, vecs, 1)
	assert.Equal(t, []float32{1, 2, 3}, vecs[0])
	assert.Equal(t, int32(2), atomic.LoadInt32(&calls), "exactly one retry after the 429")
	assert.GreaterOrEqual(t, time.Since(start), time.Second,
		"the provider must wait out the Retry-After hint before retrying")
}

// TestAPIProvider_429WithoutHintStillRetries asserts that a 429 with no
// Retry-After header still triggers exactly one retry (on a short fixed
// backoff) rather than failing immediately.
func TestAPIProvider_429WithoutHintStillRetries(t *testing.T) {
	var calls int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if atomic.AddInt32(&calls, 1) == 1 {
			w.WriteHeader(http.StatusTooManyRequests)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(openAIResponse{
			Data: []openAIEmbedding{{Embedding: []float32{4, 5, 6}, Index: 0}},
		})
	}))
	defer srv.Close()

	p := NewAPIProvider(srv.URL, "")
	vecs, err := p.EmbedBatch(context.Background(), []string{"x"})
	require.NoError(t, err)
	require.Len(t, vecs, 1)
	assert.Equal(t, int32(2), atomic.LoadInt32(&calls))
}

// TestAPIProvider_PersistentRateLimitFails asserts a server that keeps
// returning 429 eventually surfaces an error — the retry is bounded to
// one attempt, it is not an infinite loop.
func TestAPIProvider_PersistentRateLimitFails(t *testing.T) {
	var calls int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		atomic.AddInt32(&calls, 1)
		w.WriteHeader(http.StatusTooManyRequests)
	}))
	defer srv.Close()

	p := NewAPIProvider(srv.URL, "")
	_, err := p.EmbedBatch(context.Background(), []string{"x"})
	require.Error(t, err, "a persistent 429 must eventually fail")
	assert.LessOrEqual(t, atomic.LoadInt32(&calls), int32(2),
		"the 429 retry is bounded to a single re-attempt")
}

// TestAPIProvider_SendsAuthorizationHeader asserts that an embeddings API
// key (GORTEX_EMBEDDINGS_API_KEY) is forwarded as an Authorization: Bearer
// header — the fix that lets gortex use authenticated backends like OpenAI.
func TestAPIProvider_SendsAuthorizationHeader(t *testing.T) {
	t.Setenv("GORTEX_EMBEDDINGS_API_KEY", "test-secret")

	var gotAuth string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotAuth = r.Header.Get("Authorization")
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"embedding":[0.1,0.2],"index":0}]}`))
	}))
	defer srv.Close()

	// A non-Ollama URL selects the OpenAI format (the /v1/embeddings path).
	p := NewAPIProvider(srv.URL, "text-embedding-3-small")
	_, err := p.EmbedBatch(context.Background(), []string{"hello"})
	require.NoError(t, err)
	assert.Equal(t, "Bearer test-secret", gotAuth)
}

// TestAPIProvider_NoKeyOmitsAuthorizationHeader asserts that with no key
// configured, no Authorization header is sent (keyless Ollama stays keyless
// and a stray OPENAI_API_KEY does not leak to a non-OpenAI endpoint).
func TestAPIProvider_NoKeyOmitsAuthorizationHeader(t *testing.T) {
	t.Setenv("GORTEX_EMBEDDINGS_API_KEY", "")
	t.Setenv("OPENAI_API_KEY", "")

	var hadAuth bool
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, hadAuth = r.Header["Authorization"]
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"embedding":[0.1],"index":0}]}`))
	}))
	defer srv.Close()

	p := NewAPIProvider(srv.URL, "text-embedding-3-small")
	_, err := p.EmbedBatch(context.Background(), []string{"hi"})
	require.NoError(t, err)
	assert.False(t, hadAuth, "no Authorization header without a configured key")
}

// TestAPIProvider_AccumulatesTokenUsage asserts the provider reads the
// OpenAI `usage.total_tokens` field off each embedding response and
// accumulates it across calls — the signal the indexer logs so the paid
// embedding pass reports its actual token spend (it otherwise vanishes).
func TestAPIProvider_AccumulatesTokenUsage(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"embedding":[0.1,0.2],"index":0}],"usage":{"prompt_tokens":7,"total_tokens":7}}`))
	}))
	defer srv.Close()

	p := NewAPIProvider(srv.URL, "text-embedding-3-small")
	assert.Equal(t, int64(0), p.TokensUsed(), "no tokens before any call")

	_, err := p.EmbedBatch(context.Background(), []string{"hello"})
	require.NoError(t, err)
	_, err = p.EmbedBatch(context.Background(), []string{"world"})
	require.NoError(t, err)

	assert.Equal(t, int64(14), p.TokensUsed(), "usage accumulates across batches")
}

// TestAPIProvider_OpenAIBaseURLVariants asserts the OpenAI request path is
// "/v1/embeddings" whether the base URL is given with or without the "/v1"
// version segment. OpenAI-compatible bases conventionally include "/v1"
// (OpenAI, OpenRouter); without this normalization a "…/v1" base produced
// "…/v1/v1/embeddings" → 404 → silent fallback to BM25.
func TestAPIProvider_OpenAIBaseURLVariants(t *testing.T) {
	for _, suffix := range []string{"", "/v1"} {
		suffix := suffix
		t.Run("base"+suffix, func(t *testing.T) {
			var gotPath string
			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				gotPath = r.URL.Path
				w.Header().Set("Content-Type", "application/json")
				_, _ = w.Write([]byte(`{"data":[{"embedding":[0.1,0.2],"index":0}]}`))
			}))
			defer srv.Close()

			p := NewAPIProvider(srv.URL+suffix, "text-embedding-3-small")
			_, err := p.EmbedBatch(context.Background(), []string{"hi"})
			require.NoError(t, err)
			assert.Equal(t, "/v1/embeddings", gotPath,
				"both base forms must resolve to a single /v1/embeddings path")
		})
	}
}

// TestAPIProvider_ProbeDimensions asserts the probe discovers the vector
// width with exactly one embedding call, caches it (so Dimensions() then
// reports the true width up front), and is idempotent — a second probe
// issues no further HTTP. This is the fix for the daemon logging dim:0 and
// the snapshot-vector reload gate rejecting a correctly-sized cached index.
func TestAPIProvider_ProbeDimensions(t *testing.T) {
	var calls int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		atomic.AddInt32(&calls, 1)
		w.Header().Set("Content-Type", "application/json")
		// A 4-dimensional vector stands in for OpenAI's 1536-d response.
		_, _ = w.Write([]byte(`{"data":[{"embedding":[0.1,0.2,0.3,0.4],"index":0}],"usage":{"total_tokens":3}}`))
	}))
	defer srv.Close()

	p := NewAPIProvider(srv.URL, "text-embedding-3-small")
	assert.Equal(t, 0, p.Dimensions(), "width unknown before the first call")

	dim, err := p.ProbeDimensions(context.Background())
	require.NoError(t, err)
	assert.Equal(t, 4, dim, "probe reports the response vector width")
	assert.Equal(t, 4, p.Dimensions(), "width is cached after the probe")
	assert.Equal(t, int32(1), atomic.LoadInt32(&calls), "the probe is a single call")

	dim2, err := p.ProbeDimensions(context.Background())
	require.NoError(t, err)
	assert.Equal(t, 4, dim2)
	assert.Equal(t, int32(1), atomic.LoadInt32(&calls), "a second probe issues no further HTTP")
}

// TestAPIProvider_ProbeDimensionsError asserts that a probe against an
// unreachable / erroring backend surfaces the error and leaves the width at
// 0 (best-effort) — the caller only warns and indexing falls back to BM25.
func TestAPIProvider_ProbeDimensionsError(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusUnauthorized)
		_, _ = w.Write([]byte(`{"error":"bad key"}`))
	}))
	defer srv.Close()

	p := NewAPIProvider(srv.URL, "text-embedding-3-small")
	dim, err := p.ProbeDimensions(context.Background())
	require.Error(t, err, "an auth failure must surface as an error")
	assert.Equal(t, 0, dim)
	assert.Equal(t, 0, p.Dimensions(), "a failed probe leaves the width unset")
}

// TestAPIProvider_ProbeDimensionsLive hits the REAL OpenAI embeddings API to
// prove the fork's embedder is genuinely wired end-to-end (not a stub): the
// probe returns text-embedding-3-small's true 1536-d width, a batch embed
// returns 1536-d vectors, and token usage is accounted. Skipped unless a key
// is present so CI without credentials stays green.
//
//	OPENAI_API_KEY=sk-... go test ./internal/embedding -run ProbeDimensionsLive -v
func TestAPIProvider_ProbeDimensionsLive(t *testing.T) {
	if os.Getenv("GORTEX_EMBEDDINGS_API_KEY") == "" && os.Getenv("OPENAI_API_KEY") == "" {
		t.Skip("no embeddings API key (set OPENAI_API_KEY) — skipping live OpenAI probe")
	}

	p := NewAPIProvider("https://api.openai.com/v1", "text-embedding-3-small")
	dim, err := p.ProbeDimensions(context.Background())
	require.NoError(t, err, "live probe against OpenAI must succeed")
	assert.Equal(t, 1536, dim, "text-embedding-3-small is 1536-dimensional")
	assert.Equal(t, 1536, p.Dimensions())

	vecs, err := p.EmbedBatch(context.Background(), []string{"rule engine evaluates facts", "knowledge base"})
	require.NoError(t, err)
	require.Len(t, vecs, 2)
	assert.Len(t, vecs[0], 1536, "each returned vector is 1536-d")
	assert.Len(t, vecs[1], 1536)
	assert.Greater(t, p.TokensUsed(), int64(0), "the paid embedding pass accounts token usage")
}

// TestTruncateEmbedInputs asserts oversized inputs are head-truncated to the
// byte cap (so OpenAI's 8192-token limit can't abort the whole vector index)
// while normal inputs pass through untouched.
func TestTruncateEmbedInputs(t *testing.T) {
	short := "small symbol"
	long := string(make([]byte, maxEmbedInputBytes+500))

	out := truncateEmbedInputs([]string{short, long})
	assert.Equal(t, short, out[0], "short input untouched")
	assert.LessOrEqual(t, len(out[1]), maxEmbedInputBytes, "long input capped")

	in := []string{"a", "b"}
	assert.Equal(t, in, truncateEmbedInputs(in), "no oversize → same slice values")
}