ModelRouter/internal/provider/gemini_endpoint_test.go

package provider

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"sync"
	"testing"

	"llmsproxy/internal/config"
	"llmsproxy/internal/lua"
	"llmsproxy/internal/types"
)

// Gemini is the only adapter whose endpoint carries the model in the PATH
// (POST /v1beta/models/{model}:generateContent), so it is the only one that
// needs the substitution this file covers.
//
// The defect it fixes was silent: the adapter shipped endpoint="/v1/models",
// which is Gemini's model-LIST endpoint and answers 405 to POST, so the preset
// template produced a source that could never work.

func geminiVM(t *testing.T) *lua.VM {
	t.Helper()
	// The dir must be named "adapters": the VM seeds its bundled adapters into
	// <dir> on Start and a differently-named dir leaves it empty, which surfaces
	// as "adapter gemini not loaded" rather than an obvious setup error.
	vm := lua.NewVM(filepath.Join(t.TempDir(), "adapters"))
	if err := vm.Start(); err != nil {
		t.Fatalf("lua vm: %v", err)
	}
	t.Cleanup(vm.Stop)
	return vm
}

// geminiProvider builds a provider on the gemini adapter against base.
func geminiProvider(t *testing.T, base string, models ...string) *Provider {
	t.Helper()
	if len(models) == 0 {
		models = []string{"gemini-3.6-flash"}
	}
	ms := make([]config.Model, 0, len(models))
	for _, m := range models {
		ms = append(ms, config.Model{ID: m, Kind: "chat"})
	}
	return New(config.Source{
		Name:          "gem",
		BaseURL:       base,
		Adapter:       "gemini",
		APIKey:        "k",
		MaxConcurrent: 4,
		Models:        ms,
	}, geminiVM(t))
}

// geminiRecorder captures the path the gateway actually requested and replies
// with a minimal Gemini-shaped payload.
type geminiRecorder struct {
	mu     sync.Mutex
	paths  []string
	bodies []string
	stream bool
}

func (g *geminiRecorder) handler(t *testing.T) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		raw, _ := io.ReadAll(r.Body)
		g.mu.Lock()
		// RequestURI, not URL.Path: URL.Path is the DECODED path, so an escaped
		// %2F looks like a real "/" and the assertion could not tell an escaped
		// id from an unescaped one.
		g.paths = append(g.paths, r.RequestURI)
		g.bodies = append(g.bodies, string(raw))
		stream := g.stream
		g.mu.Unlock()

		if stream {
			w.Header().Set("Content-Type", "text/event-stream")
			w.WriteHeader(200)
			fmt.Fprintf(w, "data: %s\n\n", `{"candidates":[{"content":{"parts":[{"text":"hi"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":3,"candidatesTokenCount":1,"totalTokenCount":4}}`)
			fmt.Fprintln(w, "data: [DONE]")
			return
		}
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{"candidates":[{"content":{"parts":[{"text":"pong"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":3,"candidatesTokenCount":1,"totalTokenCount":4}}`)
	}
}

func (g *geminiRecorder) lastPath() string {
	g.mu.Lock()
	defer g.mu.Unlock()
	if len(g.paths) == 0 {
		return ""
	}
	return g.paths[len(g.paths)-1]
}

// TestGeminiEndpointCarriesModel: a non-streaming chat must POST to
// /v1beta/models/<model>:generateContent. The old static endpoint produced
// "/v1beta/v1/models" (a list endpoint, 405 on POST).
func TestGeminiEndpointCarriesModel(t *testing.T) {
	rec := &geminiRecorder{}
	srv := httptest.NewServer(rec.handler(t))
	defer srv.Close()

	p := geminiProvider(t, srv.URL, "gemini-3.6-flash")
	if _, err := p.Chat(context.Background(), &types.ChatRequest{
		Model:    "gemini-3.6-flash",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("hi")}},
	}); err != nil {
		t.Fatalf("chat: %v (path was %q)", err, rec.lastPath())
	}
	if got, want := rec.lastPath(), "/v1beta/models/gemini-3.6-flash:generateContent"; got != want {
		t.Errorf("path = %q, want %q", got, want)
	}
}

// TestGeminiStreamUsesStreamGenerateContent: the streaming verb lives in the
// path too, so it must be swapped for stream requests only.
func TestGeminiStreamUsesStreamGenerateContent(t *testing.T) {
	rec := &geminiRecorder{stream: true}
	srv := httptest.NewServer(rec.handler(t))
	defer srv.Close()

	p := geminiProvider(t, srv.URL, "gemini-3.6-flash")
	ch, err := p.ChatStream(context.Background(), &types.ChatRequest{
		Model:    "gemini-3.6-flash",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("hi")}},
	})
	if err != nil {
		t.Fatalf("stream: %v (path was %q)", err, rec.lastPath())
	}
	for range ch {
	}
	if got, want := rec.lastPath(), "/v1beta/models/gemini-3.6-flash:streamGenerateContent"; got != want {
		t.Errorf("path = %q, want %q", got, want)
	}
}

// TestGeminiModelFollowsRequestNotSourceDefault is the part a naive fix gets
// wrong: substituting the source's default model would make every request on a
// multi-model source bill the wrong model. AUTO pins req.Model per slot, so the
// URL must follow the REQUEST's model.
func TestGeminiModelFollowsRequestNotSourceDefault(t *testing.T) {
	rec := &geminiRecorder{}
	srv := httptest.NewServer(rec.handler(t))
	defer srv.Close()

	p := geminiProvider(t, srv.URL, "gemini-3.6-flash", "gemini-3.1-pro")
	for _, m := range []string{"gemini-3.1-pro", "gemini-3.6-flash"} {
		if _, err := p.Chat(context.Background(), &types.ChatRequest{
			Model:    m,
			Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("hi")}},
		}); err != nil {
			t.Fatalf("chat %s: %v", m, err)
		}
		if want := "/v1beta/models/" + m + ":generateContent"; rec.lastPath() != want {
			t.Errorf("path = %q, want %q", rec.lastPath(), want)
		}
	}
}

// TestGeminiModelIsEscaped: a model id reaches a URL path, so an id carrying a
// slash or space must be escaped rather than silently addressing a different
// resource (or producing a request Go's http client rejects).
func TestGeminiModelIsEscaped(t *testing.T) {
	rec := &geminiRecorder{}
	srv := httptest.NewServer(rec.handler(t))
	defer srv.Close()

	// Configure the odd id as a real model so ModelFor resolves it.
	p := geminiProvider(t, srv.URL, "vendor/model:v1 beta")
	if _, err := p.Chat(context.Background(), &types.ChatRequest{
		Model:    "vendor/model:v1 beta",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("hi")}},
	}); err != nil {
		t.Fatalf("chat: %v", err)
	}
	got := rec.lastPath()
	if strings.Contains(got, "/v1beta/models/vendor/model") {
		t.Errorf("path %q was not escaped: the id's slash changed the resource", got)
	}
	if !strings.HasPrefix(got, "/v1beta/models/") || !strings.HasSuffix(got, ":generateContent") {
		t.Errorf("path = %q, want an escaped id between the prefix and the verb", got)
	}
}

// TestNonGeminiEndpointsAreUntouched: the substitution and the verb rewrite must
// be no-ops for every adapter whose endpoint is a fixed path. A regression here
// would silently rewrite OpenAI/Anthropic/… URLs.
func TestNonGeminiEndpointsAreUntouched(t *testing.T) {
	for _, adapter := range []string{"openai", "deepseek", "anthropic", "ollama", "mistral", "trae", "sensenova"} {
		vm := geminiVM(t)
		p := New(config.Source{
			Name: "s", BaseURL: "https://example.test", Adapter: adapter,
			APIKey: "k", MaxConcurrent: 2,
			Models: []config.Model{{ID: "m1", Kind: "chat"}},
		}, vm)
		for _, stream := range []bool{false, true} {
			got := p.ChatURL("m1", stream)
			want := "https://example.test" + p.Endpoint()
			if got != want {
				t.Errorf("adapter %s stream=%v: URL = %q, want %q", adapter, stream, got, want)
			}
			if strings.Contains(got, "streamGenerateContent") {
				t.Errorf("adapter %s: streaming verb leaked into a fixed endpoint: %q", adapter, got)
			}
		}
		vm.Stop()
	}
}

// TestSourceEndpointOverridesTemplate: a source that spells its own endpoint
// must win outright — including when it spells one WITHOUT the placeholder.
func TestSourceEndpointOverridesTemplate(t *testing.T) {
	vm := geminiVM(t)
	defer vm.Stop()
	p := New(config.Source{
		Name: "s", BaseURL: "https://example.test", Adapter: "gemini",
		Endpoint: "/custom/path", APIKey: "k", MaxConcurrent: 2,
		Models: []config.Model{{ID: "m1", Kind: "chat"}},
	}, vm)
	if got, want := p.ChatURL("m1", false), "https://example.test/custom/path"; got != want {
		t.Errorf("URL = %q, want %q (source endpoint must override the adapter template)", got, want)
	}
	// And the streaming verb rewrite must not damage an unrelated path.
	if got := p.ChatURL("m1", true); strings.Contains(got, "streamGenerate") {
		t.Errorf("streaming rewrite leaked into a user-spelled path: %q", got)
	}
}

// TestGeminiResponseStillParses: the endpoint fix must not disturb the body
// transform — a request that reaches the right URL still has to come back as a
// unified response with Gemini's usageMetadata mapped.
func TestGeminiResponseStillParses(t *testing.T) {
	rec := &geminiRecorder{}
	srv := httptest.NewServer(rec.handler(t))
	defer srv.Close()

	p := geminiProvider(t, srv.URL, "gemini-3.6-flash")
	resp, err := p.Chat(context.Background(), &types.ChatRequest{
		Model:    "gemini-3.6-flash",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("hi")}},
	})
	if err != nil {
		t.Fatalf("chat: %v", err)
	}
	if resp.Content != "pong" {
		t.Errorf("content = %q, want %q", resp.Content, "pong")
	}
	if resp.TokenUsage.Prompt != 3 || resp.TokenUsage.Completion != 1 || resp.TokenUsage.Total != 4 {
		t.Errorf("usage = %+v, want prompt 3 / completion 1 / total 4", resp.TokenUsage)
	}
	// The body the adapter produced must still be Gemini-native.
	rec.mu.Lock()
	body := rec.bodies[0]
	rec.mu.Unlock()
	var sent map[string]interface{}
	if err := json.Unmarshal([]byte(body), &sent); err != nil {
		t.Fatalf("adapter sent invalid json: %v (%s)", err, body)
	}
	if _, ok := sent["contents"]; !ok {
		t.Errorf("adapter did not convert to Gemini's contents[]: %s", body)
	}
}