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()
	td := t.TempDir()
	cfgPath := filepath.Join(td, "config.yaml")
	os.WriteFile(cfgPath, []byte("listen: :0"), 0644)
	cfg := &config.Config{
		Path:        cfgPath,
		AdapterDir:  filepath.Join(td, "adapters"),
		RuntimeFile: filepath.Join(td, "runtime.json"),
		GatewayKeys: []string{"sk-test"},
		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
}

// upstreamCtrl toggles a mocked upstream's behavior between requests.
type upstreamCtrl struct {
	status int // 0 = healthy; else every request fails with that status
	hits   int // chat call count
}

// upstream returns a mocked OpenAI upstream driven by ctrl.status.
func upstream(t *testing.T, ctrl *upstreamCtrl) *httptest.Server {
	t.Helper()
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctrl.hits++
		if ctrl.status != 0 {
			w.WriteHeader(ctrl.status)
			fmt.Fprint(w, `{"error":"boom"}`)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintf(w, `{"choices":[{"message":{"content":"pong"},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4}}`)
	}))
}

// TestChatAutoChainTierFailover: AUTO chain, first slot hard-fails, the pass
// moves on within the same tier and the request is served by the next slot.
func TestChatAutoChainTierFailover(t *testing.T) {
	a, b := &upstreamCtrl{status: 500}, &upstreamCtrl{}
	aUp := upstream(t, a)
	bUp := upstream(t, b)
	defer aUp.Close()
	defer bUp.Close()
	g := newTestGateway(t,
		config.Source{Name: "a", BaseURL: aUp.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m", Priority: 100}}},
		config.Source{Name: "b", BaseURL: bUp.URL, Adapter: "openai", Models: []config.Model{{ID: "b-m", Priority: 10}}},
	)
	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 != "b-m" {
		t.Fatalf("AUTO served %q, want b-m", cc.Model)
	}
	if a.hits == 0 || b.hits == 0 {
		t.Fatalf("hit counts a=%d b=%d, want both > 0", a.hits, b.hits)
	}
}

// TestChatAutoChain503Summary: every AUTO slot fails -> 503 whose message
// names each failed tier/source/model.
func TestChatAutoChain503Summary(t *testing.T) {
	a, b := &upstreamCtrl{status: 500}, &upstreamCtrl{status: 500}
	aUp := upstream(t, a)
	bUp := upstream(t, b)
	defer aUp.Close()
	defer bUp.Close()
	g := newTestGateway(t,
		config.Source{Name: "a", BaseURL: aUp.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m", Priority: 100}}},
		config.Source{Name: "b", BaseURL: bUp.URL, Adapter: "openai", Models: []config.Model{{ID: "b-m", Priority: 10}}},
	)
	rr := doReq(t, g, "POST", "/v1/chat/completions",
		`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
	if rr.Code != http.StatusServiceUnavailable {
		t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
	}
	if !strings.Contains(rr.Body.String(), "auto providers failed") ||
		!strings.Contains(rr.Body.String(), "a/a-m") ||
		!strings.Contains(rr.Body.String(), "b/b-m") {
		t.Fatalf("503 must summarize every tier, body=%s", rr.Body.String())
	}
}

// TestChatAutoQuotaSkip: a slot whose token quota is exhausted is dropped
// from scheduling; with no other slot the chain answers 503 naming the quota.
func TestChatAutoQuotaSkip(t *testing.T) {
	ctrl := &upstreamCtrl{}
	up := upstream(t, ctrl)
	defer up.Close()
	g := newTestGateway(t,
		config.Source{Name: "a", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m", Priority: 100}}},
	)
	// one slot with an hourly quota of 1 token
	rr := doReq(t, g, "PUT", "/api/auto",
		`{"rules":[{"model":"a-m","tier":0,"token_quota":1,"period":"hour"}]}`)
	if rr.Code != 200 {
		t.Fatalf("put auto status=%d body=%s", rr.Code, rr.Body.String())
	}
	// first request consumes 4 tokens -> quota exhausted
	rr = doReq(t, g, "POST", "/v1/chat/completions",
		`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
	if rr.Code != 200 {
		t.Fatalf("first status=%d body=%s", rr.Code, rr.Body.String())
	}
	// second request must skip the exhausted slot and fail 503
	rr = doReq(t, g, "POST", "/v1/chat/completions",
		`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
	if rr.Code != http.StatusServiceUnavailable {
		t.Fatalf("quota status=%d body=%s", rr.Code, rr.Body.String())
	}
	if !strings.Contains(rr.Body.String(), "quota exhausted") {
		t.Fatalf("503 must name the quota reason, body=%s", rr.Body.String())
	}
	if ctrl.hits != 1 {
		t.Fatalf("upstream hits = %d, want 1 (exhausted slot must not be called)", ctrl.hits)
	}
}

// TestAutoStatesReportChainHealth: GET /api/auto reports per-slot health for
// the priority-page UI; a chain edit resets the failure state to zero.
func TestAutoStatesReportChainHealth(t *testing.T) {
	ctrl := &upstreamCtrl{status: 500}
	up := upstream(t, ctrl)
	defer up.Close()
	g := newTestGateway(t,
		config.Source{Name: "a", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m", Priority: 100}}},
	)
	doReq(t, g, "POST", "/v1/chat/completions",
		`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)

	fetch := func() []core.AutoSlotState {
		rr := doReq(t, g, "GET", "/api/auto", "")
		if rr.Code != 200 {
			t.Fatalf("get auto status=%d body=%s", rr.Code, rr.Body.String())
		}
		var body struct {
			Rules  []config.ModelScope  `json:"rules"`
			States []core.AutoSlotState `json:"states"`
		}
		if err := json.Unmarshal(rr.Body.Bytes(), &body); err != nil {
			t.Fatalf("unmarshal: %v", err)
		}
		return body.States
	}

	st := fetch()
	if len(st) != 1 || st[0].Model != "a-m" || st[0].Source != "a" {
		t.Fatalf("want 1 slot a/a-m, got %#v", st)
	}
	if st[0].FailCount == 0 || !st[0].Cooling {
		t.Fatalf("slot must report the failure (fail=%d cooling=%v)", st[0].FailCount, st[0].Cooling)
	}

	doReq(t, g, "PUT", "/api/auto",
		`{"rules":[{"model":"a-m","tier":0}]}`)
	st = fetch()
	if st[0].FailCount != 0 || st[0].Cooling {
		t.Fatalf("edit must reset health, got %#v", st[0])
	}
}

// TestAutoSaveResetsCooldown: editing the AUTO chain clears the cooldown of
// its slots, so a fixed upstream is schedulable again without waiting (P1).
func TestAutoSaveResetsCooldown(t *testing.T) {
	ctrl := &upstreamCtrl{status: 500}
	up := upstream(t, ctrl)
	defer up.Close()
	g := newTestGateway(t,
		config.Source{Name: "a", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m", Priority: 100}}},
	)
	rr := doReq(t, g, "POST", "/v1/chat/completions",
		`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
	if rr.Code != http.StatusServiceUnavailable {
		t.Fatalf("expect 503 while upstream down, got %d", rr.Code)
	}
	p := g.core.ProviderForSlot("a-m", "a")
	if p == nil || p.ModelAvailable("a-m") {
		t.Fatal("a-m must be cooling after the failure")
	}
	// editing the chain (same rules) must clear the cooldown immediately
	rr = doReq(t, g, "PUT", "/api/auto",
		`{"rules":[{"model":"a-m","tier":0}]}`)
	if rr.Code != 200 {
		t.Fatalf("put auto status=%d body=%s", rr.Code, rr.Body.String())
	}
	if !p.ModelAvailable("a-m") {
		t.Fatal("SaveAutoRules must reset the slot cooldown")
	}
	// healed upstream -> AUTO serves again on the next request
	ctrl.status = 0
	rr = doReq(t, g, "POST", "/v1/chat/completions",
		`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
	if rr.Code != 200 {
		t.Fatalf("AUTO after reset status=%d body=%s", rr.Code, rr.Body.String())
	}
}

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 source persisted to config.yaml
	if _, err := os.Stat(g.core.Config().Path); err != nil {
		t.Fatalf("config 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())
	}
}

func TestHasScopeModelWithSourcePrefix(t *testing.T) {
	// zen serves deepseek-v4-flash-free, so the pinning prefix strips; deepseek
	// serves deepseek-v4-pro but NOT v4-flash-free, so a model id like
	// "deepseek-v4-flash-free" (source deepseek + model id with dashes) must NOT
	// be corrupted by prefix stripping.
	g := newTestGateway(t,
		config.Source{Name: "deepseek", BaseURL: "http://d", Adapter: "openai",
			Models: []config.Model{{ID: "deepseek-v4-pro"}}},
		config.Source{Name: "zen", BaseURL: "http://z", Adapter: "openai",
			Models: []config.Model{{ID: "deepseek-v4-flash-free"}}},
	)
	scope := []config.ModelScope{{Model: "deepseek-v4-flash-free"}, {Model: "gpt-5.6-sol"}}
	for _, prefixed := range []string{"zen:deepseek-v4-flash-free", "zen/deepseek-v4-flash-free"} {
		if !g.hasScopeModel(scope, prefixed) {
			t.Errorf("hasScopeModel(scope, %q)=false, want true (prefix must be stripped)", prefixed)
		}
	}
	for _, bare := range []string{"deepseek-v4-flash-free", "gpt-5.6-sol"} {
		if !g.hasScopeModel(scope, bare) {
			t.Errorf("hasScopeModel(scope, %q)=false, want true", bare)
		}
	}
	if g.hasScopeModel(scope, "deepseek-v4-pro") {
		t.Error("hasScopeModel returned true for model outside scope")
	}
	if g.hasScopeModel(scope, "deepseek-v4-flash-free-extra") {
		t.Error("hasScopeModel returned true for unrelated model")
	}
	if !g.hasScopeModel([]config.ModelScope{{Model: "AUTO"}}, "zen:anything") {
		t.Error("AUTO scope should allow any prefixed model")
	}
}

// TestStatusRecorderFlusher pins the SSE-critical contract: the access-log
// wrapper must implement http.Flusher, otherwise the streaming chat handlers'
// w.(http.Flusher) assertion yields nil and chunks are never flushed until
// the response ends (regression for the buffered-SSE bug).
func TestStatusRecorderFlusher(t *testing.T) {
	var inner *flushRecorder
	rr := httptest.NewRecorder()
	inner = &flushRecorder{ResponseWriter: rr}
	sr := &statusRecorder{ResponseWriter: inner}

	fl, ok := any(sr).(http.Flusher)
	if !ok {
		t.Fatal("statusRecorder does not implement http.Flusher — SSE streaming is broken")
	}
	if inner.flushed {
		t.Fatal("Flush called before Flush()")
	}
	fl.Flush()
	if !inner.flushed {
		t.Fatal("statusRecorder.Flush did not forward to the underlying writer")
	}
}

// flushRecorder records whether Flush was forwarded.
type flushRecorder struct {
	http.ResponseWriter
	flushed bool
}

func (f *flushRecorder) Flush() { f.flushed = true }

// TestDirectStreamFailoverAuditSource: a direct streaming request whose first
// candidate hard-fails must be served by the fallback, and the audit record's
// Source must name the source that actually served the stream (not just the
// first candidate).
func TestDirectStreamFailoverAuditSource(t *testing.T) {
	aUp := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(500)
		fmt.Fprint(w, `{"error":"boom"}`)
	}))
	defer aUp.Close()
	bUp := mockUpstream()
	defer bUp.Close()
	g := newTestGateway(t,
		// Both sources expose the same model id so the direct scheduler has a
		// real fallback candidate after a hard-fails; b is the healthy one.
		config.Source{Name: "a", BaseURL: aUp.URL, Adapter: "openai", Models: []config.Model{{ID: "m"}}},
		config.Source{Name: "b", BaseURL: bUp.URL, Adapter: "openai", Models: []config.Model{{ID: "m"}}},
	)
	body := doReq(t, g, "POST", "/v1/chat/completions",
		`{"model":"m","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
	if body.Code != 200 {
		t.Fatalf("status=%d body=%s", body.Code, body.Body.String())
	}
	sb := body.Body.String()
	if !strings.Contains(sb, `"finish_reason":"stop"`) || !strings.Contains(sb, "[DONE]") {
		t.Fatalf("stream incomplete: %s", sb)
	}
	recs := g.stats.AuditRecords(0, 0, "")
	if len(recs) == 0 {
		t.Fatal("no audit records")
	}
	last := recs[len(recs)-1]
	if last.Source != "b" || !last.OK {
		t.Fatalf("audit rec source=%q ok=%v, want source=b ok=true (actual serving source)", last.Source, last.OK)
	}
}