ModelRouter/internal/gateway/gateway_test.go

package gateway

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

	"llmsproxy/internal/config"
	"llmsproxy/internal/core"
)

func mockUpstream() *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		var req map[string]interface{}
		_ = json.Unmarshal(body, &req)
		if stream, _ := req["stream"].(bool); stream {
			w.Header().Set("Content-Type", "text/event-stream")
			w.WriteHeader(200)
			fmt.Fprintln(w, `data: {"choices":[{"delta":{"content":"Hel"}}]}`)
			fmt.Fprintln(w, `data: {"choices":[{"delta":{"content":"lo"}}]}`)
			fmt.Fprintln(w, `data: {"choices":[{"delta":{},"finish_reason":"stop"}]}`)
			fmt.Fprintln(w, "data: [DONE]")
			return
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(200)
		fmt.Fprintf(w, `{"choices":[{"message":{"content":"pong"},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4}}`)
	}))
}

func newTestGateway(t *testing.T, srcs ...config.Source) *Gateway {
	t.Helper()
	cfg := &config.Config{
		AdapterDir:  filepath.Join(t.TempDir(), "adapters"),
		RuntimeFile: filepath.Join(t.TempDir(), "runtime.json"),
		Sources:     srcs,
	}
	if err := cfg.ApplyDefaults(); err != nil {
		t.Fatal(err)
	}
	c, err := core.NewFromConfig(cfg)
	if err != nil {
		t.Fatalf("core: %v", err)
	}
	t.Cleanup(c.Close)
	g, err := New(c, []string{"sk-test"})
	if err != nil {
		t.Fatalf("gateway: %v", err)
	}
	return g
}

func doReq(t *testing.T, g *Gateway, method, path, body string) *httptest.ResponseRecorder {
	t.Helper()
	req, _ := http.NewRequest(method, path, strings.NewReader(body))
	req.Header.Set("Authorization", "Bearer sk-test")
	if body != "" {
		req.Header.Set("Content-Type", "application/json")
	}
	rr := httptest.NewRecorder()
	g.Handler().ServeHTTP(rr, req)
	return rr
}

func TestChatSingle(t *testing.T) {
	up := mockUpstream()
	defer up.Close()
	g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
	rr := doReq(t, g, "POST", "/v1/chat/completions",
		`{"model":"mock-model","messages":[{"role":"user","content":"hi"}]}`)
	if rr.Code != 200 {
		t.Fatalf("status = %d, body=%s", rr.Code, rr.Body.String())
	}
	var cc ChatCompletion
	if err := json.Unmarshal(rr.Body.Bytes(), &cc); err != nil {
		t.Fatalf("unmarshal: %v", err)
	}
	if cc.Choices[0].Message.Content != "pong" {
		t.Fatalf("content = %q", cc.Choices[0].Message.Content)
	}
	if cc.Usage == nil || cc.Usage.Total != 4 {
		t.Fatalf("usage = %+v", cc.Usage)
	}
	if cc.Model != "mock-model" {
		t.Fatalf("model = %q", cc.Model)
	}
}

func TestChatDisableThinkingPassthrough(t *testing.T) {
	got := ""
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		b, _ := io.ReadAll(r.Body)
		got = string(b)
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{"choices":[{"message":{"content":"pong"},"finish_reason":"stop"}]}`)
	}))
	defer up.Close()
	g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
	rr := doReq(t, g, "POST", "/v1/chat/completions",
		`{"model":"mock-model","disable_thinking":true,"messages":[{"role":"user","content":"hi"}]}`)
	if rr.Code != 200 {
		t.Fatalf("status = %d, body=%s", rr.Code, rr.Body.String())
	}
	var sent map[string]interface{}
	if err := json.Unmarshal([]byte(got), &sent); err != nil {
		t.Fatalf("upstream body: %v", err)
	}
	if _, has := sent["disable_thinking"]; has {
		t.Fatalf("disable_thinking not stripped: %s", got)
	}
	// openai adapter strips disable_thinking; deepseek would map it to extra_body.thinking.
	// with the passthrough fix the flag now reaches the VM at all.
}

func TestChatMultimodalPassthrough(t *testing.T) {
	gotBody := ""
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		b, _ := io.ReadAll(r.Body)
		gotBody = string(b)
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{"choices":[{"message":{"content":"pong"},"finish_reason":"stop"}]}`)
	}))
	defer up.Close()
	g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
	rr := doReq(t, g, "POST", "/v1/chat/completions",
		`{"model":"mock-model","messages":[{"role":"user","content":[
			{"type":"text","text":"what is this?"},
			{"type":"image_url","image_url":{"url":"data:image/png;base64,QUJD"}}
		]}]}`)
	if rr.Code != 200 {
		t.Fatalf("status = %d, body=%s", rr.Code, rr.Body.String())
	}
	var sent struct {
		Messages []struct {
			Content []map[string]interface{} `json:"content"`
		} `json:"messages"`
	}
	if err := json.Unmarshal([]byte(gotBody), &sent); err != nil {
		t.Fatalf("upstream body: %v", err)
	}
	if len(sent.Messages) != 1 || len(sent.Messages[0].Content) != 2 {
		t.Fatalf("multimodal content lost: %s", gotBody)
	}
}

func TestChatAUTO(t *testing.T) {
	up := mockUpstream()
	defer up.Close()
	g := newTestGateway(t,
		config.Source{Name: "low", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "low-m", Priority: 10}}},
		config.Source{Name: "high", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "high-m", Priority: 100}}},
	)
	// no model -> AUTO -> picks the highest priority source
	rr := doReq(t, g, "POST", "/v1/chat/completions",
		`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
	if rr.Code != 200 {
		t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
	}
	var cc ChatCompletion
	_ = json.Unmarshal(rr.Body.Bytes(), &cc)
	if cc.Model != "high-m" {
		t.Fatalf("AUTO picked %q, want high-m", cc.Model)
	}
}

func TestChatAuthRequired(t *testing.T) {
	up := mockUpstream()
	defer up.Close()
	g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
	req, _ := http.NewRequest("POST", "/v1/chat/completions",
		strings.NewReader(`{"messages":[{"role":"user","content":"hi"}]}`))
	rr := httptest.NewRecorder()
	g.Handler().ServeHTTP(rr, req)
	if rr.Code != http.StatusUnauthorized {
		t.Fatalf("expected 401, got %d", rr.Code)
	}
}

func TestChatStream(t *testing.T) {
	up := mockUpstream()
	defer up.Close()
	g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
	rr := doReq(t, g, "POST", "/v1/chat/completions",
		`{"model":"mock-model","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
	body := rr.Body.String()
	if !strings.Contains(body, "data: [DONE]") {
		t.Fatalf("missing DONE, body=%s", body)
	}
	if !strings.Contains(body, "Hel") || !strings.Contains(body, "lo") {
		t.Fatalf("missing content chunks, body=%s", body)
	}
}

func TestImageGeneration(t *testing.T) {
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{"created":1,"data":[{"b64_json":"QUJD"}]}`)
	}))
	defer up.Close()
	g := newTestGateway(t, config.Source{Name: "img", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "flux-1", Kind: "image"}}})
	rr := doReq(t, g, "POST", "/v1/images/generations",
		`{"model":"flux-1","prompt":"a cat"}`)
	if rr.Code != 200 {
		t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
	}
	var out map[string]interface{}
	_ = json.Unmarshal(rr.Body.Bytes(), &out)
	data, _ := out["data"].([]interface{})
	if len(data) != 1 {
		t.Fatalf("image data len = %d", len(data))
	}
}

func TestImageAutoFallsOnlyToImageProviders(t *testing.T) {
	imageHits := 0
	chatHits := 0
	img := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		imageHits++
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{"created":1,"data":[{"b64_json":"QUJD"}]}`)
	}))
	defer img.Close()
	chatUp := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		chatHits++
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{"choices":[{"message":{"content":"hi"},"finish_reason":"stop"}]}`)
	}))
	defer chatUp.Close()
	g := newTestGateway(t,
		config.Source{Name: "chat", BaseURL: chatUp.URL, Adapter: "openai", Models: []config.Model{{ID: "chat-m", Priority: 100}}},
		config.Source{Name: "img", BaseURL: img.URL, Adapter: "openai", Models: []config.Model{{ID: "flux", Kind: "image", Priority: 1}}},
	)
	rr := doReq(t, g, "POST", "/v1/images/generations",
		`{"model":"AUTO","prompt":"a cat"}`)
	if rr.Code != 200 {
		t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
	}
	if chatHits != 0 {
		t.Fatalf("image AUTO hit chat-only provider: %d chat hits", chatHits)
	}
	if imageHits == 0 {
		t.Fatalf("image AUTO did not hit image provider")
	}
}

func TestKimicodeSigning(t *testing.T) {
	var gotAuth, gotSign string
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = io.ReadAll(r.Body)
		gotAuth = r.Header.Get("Authorization")
		gotSign = r.Header.Get("X-App-Sign")
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintf(w, `{"choices":[{"message":{"content":"ok"},"finish_reason":"stop"}]}`)
	}))
	defer up.Close()

	g := newTestGateway(t, config.Source{
		Name:    "kimi",
		BaseURL: up.URL,
		Adapter: "kimicode",
		APIKey:  "sk-kimi",
		Models:  []config.Model{{ID: "kimi-k2"}},
		Meta:    map[string]interface{}{"app_id": "app-1", "app_secret": "s3cr3t", "app_agent": "code-agent", "api_key": "sk-kimi"},
	})
	rr := doReq(t, g, "POST", "/v1/chat/completions",
		`{"model":"kimi-k2","messages":[{"role":"user","content":"hi"}]}`)
	if rr.Code != 200 {
		t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
	}
	if gotAuth == "" || !strings.Contains(gotAuth, "sk-kimi") {
		t.Fatalf("expected signed auth, got %q", gotAuth)
	}
	if gotSign == "" {
		t.Fatalf("expected app signature header")
	}
}

func TestModelRoutingPrefix(t *testing.T) {
	up := mockUpstream()
	defer up.Close()
	g := newTestGateway(t,
		config.Source{Name: "a", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "model-a"}}},
		config.Source{Name: "b", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "model-b"}}},
	)
	rr := doReq(t, g, "POST", "/v1/chat/completions",
		`{"model":"model-b","messages":[{"role":"user","content":"hi"}]}`)
	if rr.Code != 200 {
		t.Fatalf("status=%d", rr.Code)
	}
}

func TestModelsEndpoint(t *testing.T) {
	up := mockUpstream()
	defer up.Close()
	g := newTestGateway(t,
		config.Source{Name: "a", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "model-a"}}},
		config.Source{Name: "b", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "model-b"}}},
	)
	rr := doReq(t, g, "GET", "/v1/models", "")
	if rr.Code != 200 {
		t.Fatalf("status=%d", rr.Code)
	}
	var out map[string]interface{}
	_ = json.Unmarshal(rr.Body.Bytes(), &out)
	if !strings.Contains(rr.Body.String(), "model-a") || !strings.Contains(rr.Body.String(), "model-b") {
		t.Fatalf("missing models: %s", rr.Body.String())
	}
}

func TestWebUIServesPage(t *testing.T) {
	up := mockUpstream()
	defer up.Close()
	g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
	rr := doReq(t, g, "GET", "/", "")
	if rr.Code != 200 {
		t.Fatalf("status=%d", rr.Code)
	}
	if !strings.Contains(rr.Body.String(), "ModelRouter") {
		t.Fatalf("ui not served")
	}
}

func TestAdaptersAPIUpload(t *testing.T) {
	up := mockUpstream()
	defer up.Close()
	g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
	// upload adapter
	rr := doReq(t, g, "POST", "/api/adapters",
		`{"name":"testadp","code":"return {name='testadp',endpoint='/chat/completions',transform_request=function(raw) return raw end,transform_response=function(raw) return raw end}"}`)
	if rr.Code != 200 {
		t.Fatalf("upload status=%d body=%s", rr.Code, rr.Body.String())
	}
	rr = doReq(t, g, "GET", "/api/status", "")
	if !strings.Contains(rr.Body.String(), "testadp") {
		t.Fatalf("adapter not listed: %s", rr.Body.String())
	}
}

func TestSourcesAPIAddAndPersist(t *testing.T) {
	up := mockUpstream()
	defer up.Close()
	g := newTestGateway(t)
	rr := doReq(t, g, "POST", "/api/sources",
		fmt.Sprintf(`{"name":"added","base_url":"%s","adapter":"openai","models":[{"id":"new-m","priority":5}]}`, up.URL))
	if rr.Code != 200 {
		t.Fatalf("add source status=%d body=%s", rr.Code, rr.Body.String())
	}
	rr = doReq(t, g, "GET", "/v1/models", "")
	if !strings.Contains(rr.Body.String(), "new-m") {
		t.Fatalf("new model not live: %s", rr.Body.String())
	}
	// verify persistence file exists
	if _, err := os.Stat(g.core.Config().RuntimeFile); err != nil {
		t.Fatalf("runtime file not written: %v", err)
	}
}

func TestUIRequiresAuth(t *testing.T) {
	up := mockUpstream()
	defer up.Close()
	g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
	// / without a key -> redirect to /login
	req, _ := http.NewRequest("GET", "/", nil)
	rr := httptest.NewRecorder()
	g.Handler().ServeHTTP(rr, req)
	if rr.Code != http.StatusFound {
		t.Fatalf("expected 302 for /, got %d", rr.Code)
	}
	if loc := rr.Header().Get("Location"); !strings.Contains(loc, "/login") {
		t.Fatalf("expected redirect to /login, got %q", loc)
	}
	// /api/status without a key -> 401
	req, _ = http.NewRequest("GET", "/api/status", nil)
	rr = httptest.NewRecorder()
	g.Handler().ServeHTTP(rr, req)
	if rr.Code != http.StatusUnauthorized {
		t.Fatalf("expected 401 for /api/status, got %d", rr.Code)
	}
	// /login page is public
	req, _ = http.NewRequest("GET", "/login", nil)
	rr = httptest.NewRecorder()
	g.Handler().ServeHTTP(rr, req)
	if rr.Code != http.StatusOK || !strings.Contains(rr.Body.String(), "登录") {
		t.Fatalf("login page: %d %s", rr.Code, rr.Body.String())
	}
}

func TestLoginAPIAndCookie(t *testing.T) {
	up := mockUpstream()
	defer up.Close()
	g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
	// bad key -> 401
	req, _ := http.NewRequest("POST", "/api/login", strings.NewReader(`{"key":"wrong"}`))
	req.Header.Set("Content-Type", "application/json")
	rr := httptest.NewRecorder()
	g.Handler().ServeHTTP(rr, req)
	if rr.Code != http.StatusUnauthorized {
		t.Fatalf("bad login: %d", rr.Code)
	}
	// good key -> cookie
	req, _ = http.NewRequest("POST", "/api/login", strings.NewReader(`{"key":"sk-test"}`))
	req.Header.Set("Content-Type", "application/json")
	rr = httptest.NewRecorder()
	g.Handler().ServeHTTP(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("login: %d", rr.Code)
	}
	cookies := rr.Result().Cookies()
	if len(cookies) == 0 || cookies[0].Name != "gw_key" {
		t.Fatalf("no gw_key cookie set")
	}
	// use cookie to access /api/status
	req, _ = http.NewRequest("GET", "/api/status", nil)
	req.AddCookie(cookies[0])
	rr = httptest.NewRecorder()
	g.Handler().ServeHTTP(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("cookie authed status: %d", rr.Code)
	}
	var out map[string]interface{}
	_ = json.Unmarshal(rr.Body.Bytes(), &out)
	if out["base_url"] == "" {
		t.Fatalf("status missing base_url: %s", rr.Body.String())
	}
	if !strings.Contains(rr.Body.String(), "sk-test") {
		t.Fatalf("status missing gateway_keys: %s", rr.Body.String())
	}
}

func TestAPIChatInternal(t *testing.T) {
	up := mockUpstream()
	defer up.Close()
	g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
	rr := doReq(t, g, "POST", "/api/chat",
		`{"model":"mock-model","messages":[{"role":"user","content":"hi"}]}`)
	if rr.Code != 200 {
		t.Fatalf("api chat status=%d body=%s", rr.Code, rr.Body.String())
	}
	if !strings.Contains(rr.Body.String(), "pong") {
		t.Fatalf("api chat body=%s", rr.Body.String())
	}
}