ModelRouter/internal/provider/provider_test.go

package provider

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync/atomic"
	"testing"
	"time"

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

func newTestProvider(t *testing.T, src config.Source) *Provider {
	t.Helper()
	vm := lua.NewVM(filepath.Join(t.TempDir(), "adapters"))
	if err := vm.Start(); err != nil {
		t.Fatalf("vm: %v", err)
	}
	t.Cleanup(vm.Stop)
	return New(src, vm)
}

func src(name, url, adapter string, models ...string) config.Source {
	s := config.Source{Name: name, BaseURL: url, Adapter: adapter, MaxConcurrent: 4}
	for _, m := range models {
		s.Models = append(s.Models, config.Model{ID: m, Priority: 0})
	}
	return s
}

func TestProviderChat(t *testing.T) {
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var body map[string]interface{}
		_ = json.NewDecoder(r.Body).Decode(&body)
		if body["model"] != "mock-model" {
			t.Errorf("model = %v", body["model"])
		}
		fmt.Fprint(w, `{"choices":[{"message":{"content":"hi"},"finish_reason":"stop"}]}`)
	}))
	defer up.Close()
	p := newTestProvider(t, src("mock", up.URL, "openai", "mock-model"))
	resp, err := p.Chat(context.Background(), &types.ChatRequest{
		Model:    "mock-model",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("hello")}},
	})
	if err != nil {
		t.Fatalf("chat: %v", err)
	}
	if resp.Content != "hi" {
		t.Fatalf("content = %q", resp.Content)
	}
}

func TestProviderChatStream(t *testing.T) {
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n")
		fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"b\"}}]}\n\n")
		fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n")
		fmt.Fprint(w, "data: [DONE]\n\n")
	}))
	defer up.Close()
	p := newTestProvider(t, src("mock", up.URL, "openai", "m"))
	ch, err := p.ChatStream(context.Background(), &types.ChatRequest{
		Model:    "m",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
	})
	if err != nil {
		t.Fatalf("stream: %v", err)
	}
	var text string
	var done bool
	for ck := range ch {
		text += ck.Content
		if ck.Done {
			done = true
		}
	}
	if text != "ab" {
		t.Fatalf("text = %q", text)
	}
	if !done {
		t.Fatal("expected done")
	}
}

// TestStreamSuccessClearsBackoff guards P6: a clean streaming end must reset
// a previously cooled (source, model) pair.
func TestStreamSuccessClearsBackoff(t *testing.T) {
	var fail atomic.Bool
	fail.Store(true)
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if fail.Load() {
			w.WriteHeader(500)
			return
		}
		w.Header().Set("Content-Type", "text/event-stream")
		fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n")
		fmt.Fprint(w, "data: [DONE]\n\n")
	}))
	defer up.Close()
	p := newTestProvider(t, src("mock", up.URL, "openai", "m"))
	if _, err := p.Chat(context.Background(), &types.ChatRequest{
		Model:    "m",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
	}); err == nil {
		t.Fatal("expected first chat to fail")
	}
	if p.ModelAvailable("m") {
		t.Fatal("m must be cooling after the failed chat")
	}
	fail.Store(false)
	ch, err := p.ChatStream(context.Background(), &types.ChatRequest{
		Model:    "m",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
	})
	if err != nil {
		t.Fatalf("stream: %v", err)
	}
	for range ch {
	}
	if !p.ModelAvailable("m") {
		t.Fatal("clean stream must clear the cooldown")
	}
	if st := p.state("m"); st.FailCount() != 0 {
		t.Fatalf("fail count after clean stream = %d", st.FailCount())
	}
}

func TestProviderImage(t *testing.T) {
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, `{"created":123,"data":[{"b64_json":"QUJD"}]}`)
	}))
	defer up.Close()
	p := newTestProvider(t, src("img", up.URL, "openai", "flux-1"))
	resp, err := p.Image(context.Background(), &types.ImageGenRequest{Model: "flux-1", Prompt: "cat"})
	if err != nil {
		t.Fatalf("image: %v", err)
	}
	if len(resp.ImageData) != 1 || resp.ImageData[0].B64JSON != "QUJD" {
		t.Fatalf("image data = %+v", resp.ImageData)
	}
}

func TestProviderBackoffPerModel(t *testing.T) {
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(500)
		fmt.Fprint(w, "boom")
	}))
	defer up.Close()
	// two models on one source: a failure on m1 must not blacklist m2
	p := newTestProvider(t, src("mock", up.URL, "openai", "m1", "m2"))
	_, err := p.Chat(context.Background(), &types.ChatRequest{
		Model:    "m1",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
	})
	if err == nil {
		t.Fatal("expected error")
	}
	if p.ModelAvailable("m1") {
		t.Fatal("expected m1 to be cooling down")
	}
	if !p.ModelAvailable("m2") {
		t.Fatal("m2 must stay schedulable (per-model isolation)")
	}
	if !p.Available() {
		t.Fatal("source must stay available while any model is schedulable")
	}
	if st := p.state("m1"); st.FailCount() != 1 {
		t.Fatalf("fail count = %d, want 1", st.FailCount())
	}
}

func TestProviderAuthFailureSelfHeals(t *testing.T) {
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(401)
		fmt.Fprint(w, `{"error":"API_KEY_DISABLED"}`)
	}))
	defer up.Close()
	p := newTestProvider(t, src("mock2", up.URL, "openai", "m2"))
	_, err := p.Chat(context.Background(), &types.ChatRequest{
		Model:    "m2",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
	})
	if err == nil {
		t.Fatal("expected error")
	}
	st := p.state("m2")
	if st.FailCount() != backoffCapN {
		t.Fatalf("auth failure must jump to capped count, got %d", st.FailCount())
	}
	if until := st.CooldownUntil(); until <= time.Now().Add(authCooldown-time.Minute).Unix() {
		t.Fatalf("auth failure must cool for authCooldown (until=%d)", until)
	}
	if st.Pref() != -2*int64(prefFailStep) {
		t.Fatalf("auth failure pref penalty must be doubled, got %d", st.Pref())
	}
	// not permanent: the reset channel and a later success both restore it
	st.reset()
	if !p.ModelAvailable("m2") {
		t.Fatal("reset must restore schedulability")
	}
}

func TestModelStateCooldownAndRecovery(t *testing.T) {
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(500)
	}))
	defer up.Close()
	p := newTestProvider(t, src("mock", up.URL, "openai", "m"))
	_, err := p.Chat(context.Background(), &types.ChatRequest{
		Model:    "m",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
	})
	if err == nil {
		t.Fatal("expected error")
	}
	st := p.state("m")
	if st.FailCount() != 1 {
		t.Fatalf("fail count = %d", st.FailCount())
	}
	// one failure -> 5s cooldown from now
	until := st.CooldownUntil()
	want := time.Now().Add(backoffBase).Unix()
	if until < want-2 || until > want+2 {
		t.Fatalf("cooldown = %d, want ~%d", until, want)
	}
	// success resets everything and bumps the preference
	st.RecordSuccess()
	if !p.ModelAvailable("m") {
		t.Fatal("success must clear cooldown")
	}
	if st.FailCount() != 0 {
		t.Fatalf("fail count after success = %d", st.FailCount())
	}
	if st.Pref() != 1-int64(prefFailStep) {
		t.Fatalf("pref after one failure (-5) then success (+1) = %d, want %d", st.Pref(), 1-int64(prefFailStep))
	}
}

func TestModelStateRateLimitShortCooldown(t *testing.T) {
	// 429 must NOT use the exponential backoff schedule: a rate-limited but
	// quota-rich source has to re-enter rotation after the short fixed window.
	p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m"))
	st := p.state("m")

	// repeated 429s must stay at the fixed short cooldown, never escalate
	for i := 0; i < 15; i++ {
		p.ReportStatus("m", 429)
	}
	until := st.CooldownUntil()
	want := time.Now().Add(rateLimitCooldown).Unix()
	if until < want-2 || until > want+2 {
		t.Fatalf("429 cooldown = %d, want ~%d (fixed %v, not exponential)", until, want, rateLimitCooldown)
	}
	if st.FailCount() != 15 {
		t.Fatalf("fail count = %d, want 15 (counted but not escalating)", st.FailCount())
	}
	if st.Pref() != -15*int64(prefFailStep) && st.Pref() > int64(prefMin) {
		t.Fatalf("pref = %d", st.Pref())
	}
	if p.ModelAvailable("m") {
		t.Fatal("model must be cooling right after a 429")
	}
}

func TestThrottleSpacingAndCancel(t *testing.T) {
	p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m"))
	p.cfg.RPM = 120 // gap = 500ms
	p.rpmGap = time.Minute / time.Duration(p.cfg.RPM)

	start := time.Now()
	for i := 0; i < 3; i++ {
		if err := p.Throttle(context.Background()); err != nil {
			t.Fatalf("throttle %d: %v", i, err)
		}
	}
	elapsed := time.Since(start)
	// first call passes immediately, the next two wait one gap each
	want := 2 * time.Minute / time.Duration(p.cfg.RPM)
	if elapsed < want-time.Duration(100*time.Millisecond) || elapsed > want+time.Second {
		t.Fatalf("3 throttled calls took %v, want ~%v", elapsed, want)
	}

	// unlimited source: Throttle is a no-op
	p2 := newTestProvider(t, src("free", "http://127.0.0.1:1", "openai", "m"))
	if err := p2.Throttle(context.Background()); err != nil {
		t.Fatalf("unlimited throttle: %v", err)
	}

	// cancelled context aborts a pending window reservation
	p3 := newTestProvider(t, src("slow", "http://127.0.0.1:1", "openai", "m"))
	p3.cfg.RPM = 6 // 10s gap
	p3.rpmGap = time.Minute / time.Duration(p3.cfg.RPM)
	// first call passes immediately (fresh provider) and reserves the next window
	if err := p3.Throttle(context.Background()); err != nil {
		t.Fatalf("first throttle: %v", err)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	// second call must wait ~10s for the reserved window; the short deadline aborts it
	start = time.Now()
	err := p3.Throttle(ctx)
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("cancelled throttle = %v, want DeadlineExceeded", err)
	}
	if time.Since(start) > time.Second {
		t.Fatalf("cancel took %v, want fast abort", time.Since(start))
	}
}

func TestTryAcquire(t *testing.T) {
	p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m"))
	p.cfg.MaxConcurrent = 1
	p.sem = make(chan struct{}, 1)
	if err := p.TryAcquire(context.Background()); err != nil {
		t.Fatalf("first acquire: %v", err)
	}
	if err := p.TryAcquire(context.Background()); !errors.Is(err, ErrBusy) {
		t.Fatalf("second acquire = %v, want ErrBusy", err)
	}
	p.Release()
	if err := p.TryAcquire(context.Background()); err != nil {
		t.Fatalf("acquire after release: %v", err)
	}
	p.Release()
}

func TestChatBusyFailsFast(t *testing.T) {
	release := make(chan struct{})
	started := make(chan struct{}, 10)
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		started <- struct{}{}
		<-release
		fmt.Fprint(w, `{"choices":[{"message":{"content":"ok"}}]}`)
	}))
	defer up.Close()
	p := newTestProvider(t, src("mock", up.URL, "openai", "m"))
	p.cfg.MaxConcurrent = 1
	p.sem = make(chan struct{}, 1)

	done := make(chan error, 1)
	go func() {
		_, err := p.Chat(context.Background(), &types.ChatRequest{
			Model:    "m",
			Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
		})
		done <- err
	}()
	<-started // first request holds the only slot

	// second request must fail fast with ErrBusy instead of queueing
	_, err2 := p.Chat(context.Background(), &types.ChatRequest{
		Model:    "m",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
	})
	if !errors.Is(err2, ErrBusy) {
		t.Fatalf("second chat err = %v, want ErrBusy", err2)
	}
	close(release)
	if err := <-done; err != nil {
		t.Fatalf("first chat: %v", err)
	}
}

func eventually(t *testing.T, timeout time.Duration, cond func() bool, msg string) {
	t.Helper()
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatalf("timed out: %s", msg)
}

// TestChatClientCancelNotRecorded: a client disconnect before the response is
// neither success nor failure — the (source, model) state must stay clean.
func TestChatClientCancelNotRecorded(t *testing.T) {
	started := make(chan struct{})
	release := make(chan struct{})
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		close(started)
		<-release // hold the upstream open; release at cleanup
		fmt.Fprint(w, `{"choices":[{"message":{"content":"late"}}]}`)
	}))
	defer up.Close()
	p := newTestProvider(t, src("mock", up.URL, "openai", "m"))
	ctx, cancel := context.WithCancel(context.Background())
	errCh := make(chan error, 1)
	go func() {
		_, err := p.Chat(ctx, &types.ChatRequest{
			Model:    "m",
			Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
		})
		errCh <- err
	}()
	<-started
	cancel() // client disconnects mid-request
	if err := <-errCh; err == nil {
		t.Fatal("cancelled request must return an error")
	}
	close(release)
	eventually(t, 2*time.Second, func() bool {
		_, fc, _ := p.ModelHealthInfo("m")
		return fc == 0
	}, "client cancel must not record a scheduling failure")
}

// TestChatStreamEmptyBodyNotSuccess: a 200 that yields zero chunks and no
// [DONE] must fail the candidate (error return, no delivered chunks) so the
// scheduler falls through to the next source, and record the failure.
func TestChatStreamEmptyBodyNotSuccess(t *testing.T) {
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		// 200 with an empty body: no chunks, no [DONE]
	}))
	defer up.Close()
	p := newTestProvider(t, src("mock", up.URL, "openai", "m"))
	ch, err := p.ChatStream(context.Background(), &types.ChatRequest{
		Model:    "m",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
	})
	if err == nil {
		if ch != nil {
			for range ch {
			}
		}
		t.Fatal("empty stream must fail the candidate for scheduler fallback")
	}
	eventually(t, 2*time.Second, func() bool {
		_, fc, _ := p.ModelHealthInfo("m")
		return fc >= 1
	}, "empty stream must record a failure")
}

// TestChatStreamErrorFinishFailsCandidate: zen free pool answers HTTP 200
// with a single-chunk stream whose only payload is finish_reason:"network_error"
// and an empty delta. That must fail the candidate (not deliver a laundered
// empty reply), so the scheduler can fall through to the next source.
func TestChatStreamErrorFinishFailsCandidate(t *testing.T) {
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		fmt.Fprint(w, `data: {"choices":[{"index":0,"finish_reason":"network_error","delta":{"role":"assistant","content":""}}]}`+"\n\n")
		fmt.Fprint(w, "data: [DONE]\n\n")
	}))
	defer up.Close()
	p := newTestProvider(t, src("mock", up.URL, "opencode", "m"))
	_, err := p.ChatStream(context.Background(), &types.ChatRequest{
		Model:    "m",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
	})
	if err == nil {
		t.Fatal("network_error-only stream must fail the candidate")
	}
	eventually(t, 2*time.Second, func() bool {
		_, fc, _ := p.ModelHealthInfo("m")
		return fc >= 1
	}, "error-only stream must record a failure")
}

// TestChatStreamInstantEmptyStopStillDelivered: a legitimate completion that
// ends immediately with the standard finish_reason:"stop" and zero content is
// NOT an upstream error and must still reach the client.
func TestChatStreamInstantEmptyStopStillDelivered(t *testing.T) {
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		fmt.Fprint(w, `data: {"choices":[{"index":0,"finish_reason":"stop","delta":{}}]}`+"\n\n")
		fmt.Fprint(w, "data: [DONE]\n\n")
	}))
	defer up.Close()
	p := newTestProvider(t, src("mock", up.URL, "openai", "m"))
	ch, err := p.ChatStream(context.Background(), &types.ChatRequest{
		Model:    "m",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
	})
	if err != nil {
		t.Fatalf("standard stop stream must be delivered: %v", err)
	}
	n := 0
	for range ch {
		n++
	}
	if n == 0 {
		t.Fatal("expected at least the terminating chunk")
	}
	eventually(t, 2*time.Second, func() bool {
		_, fc, _ := p.ModelHealthInfo("m")
		return fc == 0
	}, "standard stop must count as success")
}

// TestImageAutoUsesImageModel: AUTO image generation on a mixed source must
// send the image-kind model id, never the best chat model.
func TestImageAutoUsesImageModel(t *testing.T) {
	var gotModel string
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var body map[string]interface{}
		_ = json.NewDecoder(r.Body).Decode(&body)
		gotModel, _ = body["model"].(string)
		fmt.Fprint(w, `{"created":1,"data":[{"url":"http://x/1.png"}]}`)
	}))
	defer up.Close()
	s := config.Source{Name: "mix", BaseURL: up.URL, Adapter: "openai", MaxConcurrent: 4}
	s.Models = []config.Model{
		{ID: "chat-m", Kind: "chat", Priority: 100},
		{ID: "img-m", Kind: "image", Priority: 50},
	}
	p := newTestProvider(t, s)
	resp, err := p.Image(context.Background(), &types.ImageGenRequest{Prompt: "cat", Model: "AUTO"})
	if err != nil {
		t.Fatalf("image: %v", err)
	}
	if len(resp.ImageData) == 0 {
		t.Fatal("no image data returned")
	}
	if gotModel != "img-m" {
		t.Fatalf("AUTO image must use the image-kind model, got %q", gotModel)
	}
}

func TestStandardSSEChunkFinishReason(t *testing.T) {
	out := standardSSEChunk(`{"choices":[{"index":0,"finish_reason":"tool_calls","delta":{}}]}`)
	if !strings.Contains(out, `"done":true`) || !strings.Contains(out, `"finish_reason":"tool_calls"`) {
		t.Fatalf("real finish reason must pass through: %s", out)
	}
	out = standardSSEChunk(`{"choices":[{"index":0,"finish_reason":"","delta":{"content":"hi"}}]}`)
	if strings.Contains(out, `"Done":true`) {
		t.Fatalf("empty-string finish reason is not a finish signal: %s", out)
	}
}

// TestAdapterHookCondensesError: the adapter transform_error hook owns the
// per-source error format; its reason must reach the client verbatim.
func TestAdapterHookCondensesError(t *testing.T) {
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(503)
		fmt.Fprint(w, `{"error":{"type":"FreeUsageLimitError","message":"Rate limit exceeded. Please try again later."}}`)
	}))
	defer up.Close()
	p := newTestProvider(t, src("mock", up.URL, "opencode", "m"))
	_, err := p.Chat(context.Background(), &types.ChatRequest{
		Model:    "m",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
	})
	if err == nil || !strings.Contains(err.Error(), "zen free pool quota exhausted") {
		t.Fatalf("adapter hook reason must surface, got %v", err)
	}
	if strings.Contains(err.Error(), "Rate limit exceeded") {
		t.Fatalf("raw upstream body must not leak past the hook: %v", err)
	}
}

// TestUnknownErrorFallbackWithoutHook: an adapter without transform_error
// gets the uniform core fallback; the raw body must not leak to clients.
func TestUnknownErrorFallbackWithoutHook(t *testing.T) {
	dir := filepath.Join(t.TempDir(), "adapters")
	if err := os.MkdirAll(dir, 0o755); err != nil {
		t.Fatal(err)
	}
	minimal := `return { name="custom", endpoint="/chat/completions",
	transform_request=function(raw) return raw end,
	transform_response=function(raw) return raw end }
`
	if err := os.WriteFile(filepath.Join(dir, "custom.lua"), []byte(minimal), 0o644); err != nil {
		t.Fatal(err)
	}
	vm := lua.NewVM(dir)
	if err := vm.Start(); err != nil {
		t.Fatalf("vm: %v", err)
	}
	defer vm.Stop()
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(500)
		fmt.Fprint(w, `{"weird":{"shape":["no","message"]}}`)
	}))
	defer up.Close()
	p := New(src("mock", up.URL, "custom", "m"), vm)
	_, err := p.Chat(context.Background(), &types.ChatRequest{
		Model:    "m",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
	})
	if err == nil || !strings.Contains(err.Error(), "unknown error") {
		t.Fatalf("hook-less adapter must fall back to unknown error, got %v", err)
	}
	if strings.Contains(err.Error(), "shape") {
		t.Fatalf("raw body must not leak: %v", err)
	}
}

// TestToolIndexRemapper verifies sparse upstream tool-call indices are
// compacted to dense 0-based OpenAI ordinals (Issue 2). Anthropic emits the
// content_block ordinal, so a thinking block at index 0 pushes the first
// tool_use to index 1 — OpenAI clients accumulating by index would leave a
// hole at 0 and mis-assemble multi-tool responses.
func TestToolIndexRemapper(t *testing.T) {
	t.Run("single tool after thinking block", func(t *testing.T) {
		r := newToolIndexRemapper()
		// content_block index 1 (thinking was 0) -> dense 0
		out := r.remap(json.RawMessage(`[{"index":1,"id":"toolu_a","type":"function","function":{"name":"calc","arguments":""}}]`))
		var got []map[string]interface{}
		if err := json.Unmarshal(out, &got); err != nil {
			t.Fatalf("unmarshal: %v (%s)", err, out)
		}
		if got[0]["index"] != float64(0) {
			t.Fatalf("index = %v, want 0: %s", got[0]["index"], out)
		}
		// id/name/arguments must survive verbatim
		if got[0]["id"] != "toolu_a" {
			t.Fatalf("id lost: %s", out)
		}
		fn := got[0]["function"].(map[string]interface{})
		if fn["name"] != "calc" {
			t.Fatalf("function.name lost: %s", out)
		}
		// subsequent argument fragments on the SAME upstream index reuse ordinal 0
		out2 := r.remap(json.RawMessage(`[{"index":1,"id":"","type":"function","function":{"name":"","arguments":"{\"a\":1}"}}]`))
		var got2 []map[string]interface{}
		if err := json.Unmarshal(out2, &got2); err != nil {
			t.Fatalf("unmarshal2: %v", err)
		}
		if got2[0]["index"] != float64(0) {
			t.Fatalf("fragment index = %v, want stable 0: %s", got2[0]["index"], out2)
		}
	})

	t.Run("multiple tools get distinct dense ordinals", func(t *testing.T) {
		r := newToolIndexRemapper()
		// thinking=0, tool_use=1, tool_use=2  ->  0, 1
		a := r.remap(json.RawMessage(`[{"index":1,"id":"t1","type":"function","function":{"name":"f1","arguments":""}}]`))
		b := r.remap(json.RawMessage(`[{"index":2,"id":"t2","type":"function","function":{"name":"f2","arguments":""}}]`))
		var ga, gb []map[string]interface{}
		json.Unmarshal(a, &ga)
		json.Unmarshal(b, &gb)
		if ga[0]["index"] != float64(0) {
			t.Fatalf("first tool index = %v, want 0", ga[0]["index"])
		}
		if gb[0]["index"] != float64(1) {
			t.Fatalf("second tool index = %v, want 1", gb[0]["index"])
		}
		// interleaved fragments keep their own ordinals
		a2 := r.remap(json.RawMessage(`[{"index":1,"id":"","type":"function","function":{"name":"","arguments":"x"}}]`))
		var ga2 []map[string]interface{}
		json.Unmarshal(a2, &ga2)
		if ga2[0]["index"] != float64(0) {
			t.Fatalf("t1 fragment index = %v, want 0", ga2[0]["index"])
		}
	})

	t.Run("already dense indices pass through untouched", func(t *testing.T) {
		r := newToolIndexRemapper()
		in := json.RawMessage(`[{"index":0,"id":"t","type":"function","function":{"name":"f","arguments":""}}]`)
		out := r.remap(in)
		if string(out) != string(in) {
			t.Fatalf("dense input was rewritten:\n in=%s\nout=%s", in, out)
		}
	})

	t.Run("empty and malformed input is safe", func(t *testing.T) {
		r := newToolIndexRemapper()
		if got := r.remap(nil); got != nil {
			t.Fatalf("nil -> %s", got)
		}
		bad := json.RawMessage(`not json`)
		if got := r.remap(bad); string(got) != string(bad) {
			t.Fatalf("malformed input must pass through: %s", got)
		}
	})
}

// ---- half-cooldown probe gate (plan 阶段 1) ----

// TestProbeGateHalfCooldown pins the core rule: a cooling model is fully silent
// in the first half of its cooldown window and probeable exactly once in the
// second half.
func TestProbeGateHalfCooldown(t *testing.T) {
	var st ModelState
	// open a 100s window centred so that "now" sits in the first half
	now := time.Now().Unix()
	st.cooldownFrom.Store(now - 10)
	st.cooldownUntil.Store(now + 90)
	if st.TryProbe() {
		t.Fatal("first half of the cooldown window must stay silent")
	}

	// move into the second half
	st.cooldownFrom.Store(now - 90)
	st.cooldownUntil.Store(now + 10)
	if !st.TryProbe() {
		t.Fatal("second half of the cooldown window must allow one probe")
	}
	if st.TryProbe() {
		t.Fatal("probe permit must be single-token")
	}
	if !st.Probing() {
		t.Fatal("Probing() must report the claimed permit")
	}
	st.ProbeDone()
	if st.Probing() {
		t.Fatal("ProbeDone must release the permit")
	}
	if !st.TryProbe() {
		t.Fatal("permit must be reusable after ProbeDone")
	}
}

// TestProbeGateNoProbeWhenHealthy: a healthy model needs no probe permit — the
// caller schedules it normally.
func TestProbeGateNoProbeWhenHealthy(t *testing.T) {
	var idle ModelState
	if idle.TryProbe() {
		t.Fatal("a healthy model must not be probed (schedule it normally)")
	}
}

// TestProbeGateRescuesPrefFloor documents a fix that comes with the probe gate:
// a slot whose preference sank to prefMin used to be refused by ModelAvailable
// forever, even long after every cooldown expired — a permanent blacklist. It is
// now probeable, so one success brings it back.
func TestProbeGateRescuesPrefFloor(t *testing.T) {
	var dead ModelState
	dead.pref.Store(prefMin)
	if !dead.TryProbe() {
		t.Fatal("a slot stuck at prefMin must be probeable, not blacklisted")
	}
	dead.RecordSuccess()
	dead.ProbeDone()
	if dead.Pref() <= prefMin {
		t.Fatalf("one probe success must lift the slot off the floor, pref=%d", dead.Pref())
	}
}

// TestProbeSuccessRestoresImmediately is the user-visible payoff: a recovered
// upstream must not sit out the rest of its cooldown.
func TestProbeSuccessRestoresImmediately(t *testing.T) {
	var st ModelState
	for i := 0; i < backoffCapN; i++ {
		st.RecordFailure(false)
	}
	if cd := st.CooldownUntil() - time.Now().Unix(); cd < int64(backoffCap.Seconds())-2 {
		t.Fatalf("10 failures must cool for backoffCap, got %ds", cd)
	}
	// jump into the probe window: the midpoint of [from, until) must be in the
	// past, so push the window's start back past its full length
	now := time.Now().Unix()
	st.cooldownFrom.Store(now - int64(backoffCap.Seconds()) - 10)
	if !st.TryProbe() {
		t.Fatal("probe must be allowed past the midpoint")
	}
	if !st.Probing() {
		t.Fatal("probe permit must be held while the probe runs")
	}
	st.RecordSuccess()
	st.ProbeDone()
	if !st.Available() {
		t.Fatal("probe success must clear the cooldown immediately")
	}
	if st.FailCount() != 0 {
		t.Fatalf("probe success must reset failCount, got %d", st.FailCount())
	}
	if st.CooldownUntil() != 0 || st.CooldownFrom() != 0 {
		t.Fatal("probe success must clear both cooldown bounds")
	}
}

// TestProbeFailureDefersNextProbe: a failed probe must not immediately retry —
// the new window's midpoint pushes the next probe out.
func TestProbeFailureDefersNextProbe(t *testing.T) {
	var st ModelState
	st.RecordFailure(false) // 5s window
	now := time.Now().Unix()
	st.cooldownFrom.Store(now - 100)
	st.cooldownUntil.Store(now + 10)
	if !st.TryProbe() {
		t.Fatal("probe must be allowed past the midpoint")
	}
	st.RecordFailure(false) // probe failed: a fresh, longer window opens
	st.ProbeDone()
	if st.TryProbe() {
		t.Fatal("a failed probe must defer the next probe to the new midpoint")
	}
	if from := st.CooldownFrom(); from < now {
		t.Fatalf("failed probe must reopen the window (from=%d, now=%d)", from, now)
	}
}

// TestAuthCooldownIsSeparateFromCap documents that credential failures cool on
// their own (longer) schedule instead of reusing the exponential cap.
func TestAuthCooldownIsSeparateFromCap(t *testing.T) {
	var st ModelState
	st.RecordFailure(true)
	cd := st.CooldownUntil() - time.Now().Unix()
	if cd < int64(authCooldown.Seconds())-2 || cd > int64(authCooldown.Seconds())+2 {
		t.Fatalf("auth failure must cool for authCooldown (%v), got %ds", authCooldown, cd)
	}
	if st.FailCount() != backoffCapN {
		t.Fatalf("auth failure must persist the capped count, got %d", st.FailCount())
	}
}

// TestQuotaExhaustedDoesNotEscalate is the second half of the user's complaint:
// running out of allowance must not be treated as a failure streak.
func TestQuotaExhaustedDoesNotEscalate(t *testing.T) {
	var st ModelState
	st.RecordQuotaExhausted(2 * time.Minute)
	if st.FailCount() != 0 {
		t.Fatalf("quota exhaustion must not bump failCount, got %d", st.FailCount())
	}
	cd := st.CooldownUntil() - time.Now().Unix()
	if cd < 118 || cd > 122 {
		t.Fatalf("quota cooldown must honour the reset window, got %ds", cd)
	}

	// an unbounded window is clamped so a topped-up allowance is still noticed
	var big ModelState
	big.RecordQuotaExhausted(72 * time.Hour)
	if cd := big.CooldownUntil() - time.Now().Unix(); cd > int64(maxQuotaCooldown.Seconds())+2 {
		t.Fatalf("quota cooldown must be capped at maxQuotaCooldown, got %ds", cd)
	}

	// no hint at all falls back to the short window
	var none ModelState
	none.RecordQuotaExhausted(0)
	if cd := none.CooldownUntil() - time.Now().Unix(); cd > int64(rateLimitCooldown.Seconds())+2 {
		t.Fatalf("unknown reset window must fall back to rateLimitCooldown, got %ds", cd)
	}
}

// TestQuotaReasonRouting checks that a 429 carrying a quota phrase is cooled as
// quota exhaustion while a plain rate limit keeps the short fixed window.
func TestQuotaReasonRouting(t *testing.T) {
	p := newTestProvider(t, src("q", "http://127.0.0.1:1", "openai", "m"))
	p.ReportStatusReason("m", 429, "429 insufficient_quota: you exceeded your current quota")
	st := p.state("m")
	if st.FailCount() != 0 {
		t.Fatalf("quota rejection must not bump failCount, got %d", st.FailCount())
	}
	if cd := st.CooldownUntil() - time.Now().Unix(); cd <= int64(rateLimitCooldown.Seconds()) {
		t.Fatalf("quota rejection must cool longer than a plain 429, got %ds", cd)
	}

	p2 := newTestProvider(t, src("r", "http://127.0.0.1:1", "openai", "m"))
	p2.ReportStatusReason("m", 429, "rate limit reached, slow down")
	st2 := p2.state("m")
	if cd := st2.CooldownUntil() - time.Now().Unix(); cd > int64(rateLimitCooldown.Seconds())+2 {
		t.Fatalf("plain 429 must keep the short window, got %ds", cd)
	}
}

// TestModelSchedulableProbeHandoff covers the provider-level gate the scheduler
// consumes.
func TestModelSchedulableProbeHandoff(t *testing.T) {
	p := newTestProvider(t, src("s", "http://127.0.0.1:1", "openai", "m"))
	if ok, probe := p.ModelSchedulable("m"); !ok || probe {
		t.Fatalf("healthy model must be normally schedulable, got ok=%v probe=%v", ok, probe)
	}
	st := p.state("m")
	st.RecordFailure(false)
	now := time.Now().Unix()
	st.cooldownFrom.Store(now - 100)
	st.cooldownUntil.Store(now + 10)

	ok, probe := p.ModelSchedulable("m")
	if !ok || !probe {
		t.Fatalf("cooling model past midpoint must be probeable, got ok=%v probe=%v", ok, probe)
	}
	if ok2, _ := p.ModelSchedulable("m"); ok2 {
		t.Fatal("only one probe permit may be outstanding")
	}
	p.ProbeDone("m")
	if ok3, probe3 := p.ModelSchedulable("m"); !ok3 || !probe3 {
		t.Fatal("permit must be reclaimable after ProbeDone")
	}
}

// TestProbeSlotsFollowsMaxConcurrent pins the user's example: max_concurrent=10
// yields exactly one probe slot.
func TestProbeSlotsFollowsMaxConcurrent(t *testing.T) {
	for _, tc := range []struct{ mc, want int }{
		{0, 1}, {1, 1}, {4, 1}, {10, 1}, {20, 2}, {100, 2},
	} {
		s := src("p", "http://127.0.0.1:1", "openai", "m")
		s.MaxConcurrent = tc.mc
		p := newTestProvider(t, s)
		if got := p.ProbeSlots(); got != tc.want {
			t.Fatalf("max_concurrent=%d: probe slots = %d, want %d", tc.mc, got, tc.want)
		}
	}
}

// TestModelProbeInfo checks the UI-facing view of the probe window.
func TestModelProbeInfo(t *testing.T) {
	p := newTestProvider(t, src("i", "http://127.0.0.1:1", "openai", "m"))
	if from, after, probing := p.ModelProbeInfo("m"); from != 0 || after != 0 || probing {
		t.Fatalf("idle model must report no probe window, got %d %d %v", from, after, probing)
	}
	st := p.state("m")
	now := time.Now().Unix()
	st.cooldownFrom.Store(now)
	st.cooldownUntil.Store(now + 100)
	from, after, _ := p.ModelProbeInfo("m")
	if from != now || after != now+50 {
		t.Fatalf("probe window = (%d,%d), want (%d,%d)", from, after, now, now+50)
	}
}

// TestUpstreamRecoveryWithoutWaitingOutCooldown is the end-to-end shape of the
// user's complaint: an upstream that 500s ten times lands in the capped
// cooldown, and used to be unusable for a full 30 minutes even after it
// recovered. With the probe gate the slot is silent for the first half of the
// window and then serves the very next request itself, clearing the cooldown.
//
// Real time is not waited out: the window bounds are moved so the test observes
// the same state the scheduler would see at cap/2.
func TestUpstreamRecoveryWithoutWaitingOutCooldown(t *testing.T) {
	var healthy atomic.Bool
	var hits atomic.Int64
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		if !healthy.Load() {
			w.WriteHeader(500)
			return
		}
		fmt.Fprint(w, `{"choices":[{"message":{"content":"back"},"finish_reason":"stop"}]}`)
	}))
	defer up.Close()
	p := newTestProvider(t, src("recov", up.URL, "openai", "m"))
	req := func() *types.ChatRequest {
		return &types.ChatRequest{
			Model:    "m",
			Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
		}
	}

	// 1. drive the slot into the capped cooldown
	for i := 0; i < backoffCapN; i++ {
		if _, err := p.Chat(context.Background(), req()); err == nil {
			t.Fatalf("attempt %d: expected upstream failure", i)
		}
	}
	st := p.state("m")
	if p.ModelAvailable("m") {
		t.Fatal("10 consecutive 5xx must take the slot out of normal rotation")
	}
	cd := st.CooldownUntil() - time.Now().Unix()
	if cd < int64(backoffCap.Seconds())-2 {
		t.Fatalf("cooldown = %ds, want ~%v", cd, backoffCap)
	}
	// the old 30-minute lockout is gone
	if cd > int64((6 * time.Minute).Seconds()) {
		t.Fatalf("cooldown %ds is too long: the cap must be minutes, not half an hour", cd)
	}

	// 2. upstream recovers, but we are still in the FIRST half of the window:
	//    the slot must stay completely silent (no probe traffic at all).
	healthy.Store(true)
	quiet := hits.Load()
	if ok, probe := p.ModelSchedulable("m"); ok || probe {
		t.Fatalf("first half of the window must stay silent, got ok=%v probe=%v", ok, probe)
	}
	if hits.Load() != quiet {
		t.Fatal("no request may reach the upstream during the silent half")
	}

	// 3. advance past the midpoint: exactly one probe is admitted.
	st.cooldownFrom.Store(time.Now().Unix() - int64(backoffCap.Seconds()) - 1)
	ok, isProbe := p.ModelSchedulable("m")
	if !ok || !isProbe {
		t.Fatalf("past the midpoint one probe must be admitted, got ok=%v probe=%v", ok, isProbe)
	}
	if ok2, _ := p.ModelSchedulable("m"); ok2 {
		t.Fatal("only ONE probe may be in flight while cooling")
	}

	// 4. the probe is a real request: its success restores the slot fully.
	resp, err := p.Chat(context.Background(), req())
	p.ProbeDone("m")
	if err != nil {
		t.Fatalf("probe request failed: %v", err)
	}
	if resp.Content != "back" {
		t.Fatalf("probe response = %q, want \"back\"", resp.Content)
	}
	if !p.ModelAvailable("m") {
		t.Fatal("a successful probe must return the slot to normal rotation immediately")
	}
	if st.CooldownUntil() != 0 || st.FailCount() != 0 {
		t.Fatalf("probe success must clear cooldown/failCount, got until=%d fails=%d",
			st.CooldownUntil(), st.FailCount())
	}
	if ok, probe := p.ModelSchedulable("m"); !ok || probe {
		t.Fatalf("recovered slot must schedule normally, got ok=%v probe=%v", ok, probe)
	}
}

// TestRecoveryProbeBudgetIsBounded quantifies the cost of the shorter cap
// against a still-dead upstream: one probe per window, not a retry storm.
func TestRecoveryProbeBudgetIsBounded(t *testing.T) {
	var hits atomic.Int64
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		w.WriteHeader(500)
	}))
	defer up.Close()
	p := newTestProvider(t, src("dead", up.URL, "openai", "m"))
	req := &types.ChatRequest{
		Model:    "m",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
	}
	for i := 0; i < backoffCapN; i++ {
		_, _ = p.Chat(context.Background(), req)
	}
	baseline := hits.Load()
	st := p.state("m")

	// simulate three consecutive probe windows against a dead upstream
	for i := 0; i < 3; i++ {
		st.cooldownFrom.Store(time.Now().Unix() - int64(backoffCap.Seconds()) - 1)
		ok, isProbe := p.ModelSchedulable("m")
		if !ok || !isProbe {
			t.Fatalf("window %d: expected a probe permit", i)
		}
		// a second caller in the same window gets nothing
		if ok2, _ := p.ModelSchedulable("m"); ok2 {
			t.Fatalf("window %d: more than one probe admitted", i)
		}
		if _, err := p.Chat(context.Background(), req); err == nil {
			t.Fatalf("window %d: expected the probe to fail", i)
		}
		p.ProbeDone("m")
	}
	if extra := hits.Load() - baseline; extra != 3 {
		t.Fatalf("3 probe windows must cost exactly 3 upstream requests, got %d", extra)
	}
}