openteam/server/internal/proxy/hold_test.go

package proxy

import (
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/openteam/server/internal/pkg/crypto"
	"github.com/openteam/server/internal/store"
)

// addChannel 给 harness 加一个绑定同一模型的渠道，返回渠道名。
// weight 传大数可以把加权随机的起点压到该渠道上（与既有测试同一手法）。
func addChannel(t *testing.T, h *harness, name, baseURL string, weight int) {
	t.Helper()
	encKey, err := crypto.NewEncryptor("test-master-key-1234567890").Encrypt("upstream-key")
	if err != nil {
		t.Fatal(err)
	}
	var m store.Model
	h.db.Where("name = ?", "m1").First(&m)
	ch := store.Channel{Name: name, Provider: store.ChannelProviderOpenAI,
		Formats: []string{store.FormatChat}, BaseURL: baseURL, APIKeyEnc: encKey,
		Weight: weight, TimeoutMS: 0, MaxConcurrency: 4,
		HealthStatus: store.ChannelHealthHealthy, Enabled: true}
	if err := h.db.Create(&ch).Error; err != nil {
		t.Fatal(err)
	}
	if err := h.db.Create(&store.ChannelModelBinding{ChannelID: ch.ID, ModelID: m.ID, UpstreamModel: "m1"}).Error; err != nil {
		t.Fatal(err)
	}
}

// pushWeightTo 把指定名字渠道的权重改大，让加权随机几乎必然从它开始。
func pushWeightTo(t *testing.T, h *harness, name string, weight int) {
	t.Helper()
	if err := h.db.Model(&store.Channel{}).Where("name = ?", name).Update("weight", weight).Error; err != nil {
		t.Fatal(err)
	}
}

// 流式容错：渠道1 在首个内容事件之前回错误事件（200+流式形式的错误），应换渠道2
// 重试——客户端只看到渠道2 的内容，全程 200，无感知切换。
func TestStreamFailoverBeforeFirstContent(t *testing.T) {
	var ch1Hits, ch2Hits atomic.Int64
	bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ch1Hits.Add(1)
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		_, _ = io.WriteString(w, `data: {"error":{"message":"overloaded","type":"server_error"}}`+"\n\n")
	}))
	defer bad.Close()
	good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ch2Hits.Add(1)
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		_, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"来自健康渠道"}}]}`+"\n\n")
		_, _ = io.WriteString(w, "data: [DONE]\n\n")
	}))
	defer good.Close()

	h := newHarness(t, bad.URL, 0, true, nil)
	addChannel(t, h, "c2", good.URL, 1)
	pushWeightTo(t, h, "c1", 1000)

	w := h.run(t, nil, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
	if w.Code != http.StatusOK {
		t.Fatalf("code=%d body=%s", w.Code, w.Body.String())
	}
	out := w.Body.String()
	if ch1Hits.Load() != 1 || ch2Hits.Load() != 1 {
		t.Fatalf("渠道命中次数不符：ch1=%d ch2=%d", ch1Hits.Load(), ch2Hits.Load())
	}
	if strings.Contains(out, "overloaded") {
		t.Fatalf("渠道1 的错误不应到达客户端：%q", out)
	}
	if !strings.Contains(out, "来自健康渠道") {
		t.Fatalf("应收到渠道2 的内容：%q", out)
	}
}

// Anthropic 形状的错误事件（{"type":"error","error":{...}}）同样在放行前被识别换渠道。
func TestAnthropicStyleStreamErrorFailsOver(t *testing.T) {
	bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		_, _ = io.WriteString(w, `event: error`+"\n"+`data: {"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}`+"\n\n")
	}))
	defer bad.Close()
	good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		_, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"ok"}}]}`+"\n\n")
		_, _ = io.WriteString(w, "data: [DONE]\n\n")
	}))
	defer good.Close()

	h := newHarness(t, bad.URL, 0, true, nil)
	addChannel(t, h, "c2", good.URL, 1)
	pushWeightTo(t, h, "c1", 1000)

	w := h.run(t, nil, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
	if w.Code != http.StatusOK || !strings.Contains(w.Body.String(), `"content":"ok"`) {
		t.Fatalf("应无感知切换到渠道2：code=%d body=%s", w.Code, w.Body.String())
	}
}

// 渠道回了 200+event-stream 却一个字没写就断流：放行前视为失败，换渠道重试。
func TestEmptyStreamBeforeContentFailsOver(t *testing.T) {
	bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		w.(http.Flusher).Flush()
		// 直接结束：EOF 且无内容
	}))
	defer bad.Close()
	good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		_, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"补位渠道"}}]}`+"\n\n")
		_, _ = io.WriteString(w, "data: [DONE]\n\n")
	}))
	defer good.Close()

	h := newHarness(t, bad.URL, 0, true, nil)
	addChannel(t, h, "c2", good.URL, 1)
	pushWeightTo(t, h, "c1", 1000)

	w := h.run(t, nil, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
	if w.Code != http.StatusOK || !strings.Contains(w.Body.String(), "补位渠道") {
		t.Fatalf("空流应触发故障转移：code=%d body=%s", w.Code, w.Body.String())
	}
}

// didNotStream：要流式，渠道却回了整包 JSON——视为失败换渠道，而不是把 JSON 塞给等 SSE 的客户端。
func TestDidNotStreamFailsOver(t *testing.T) {
	bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, _ = io.WriteString(w, `{"id":"x","choices":[{"message":{"role":"assistant","content":"整包"}}]}`)
	}))
	defer bad.Close()
	good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		_, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"真流式"}}]}`+"\n\n")
		_, _ = io.WriteString(w, "data: [DONE]\n\n")
	}))
	defer good.Close()

	h := newHarness(t, bad.URL, 0, true, nil)
	addChannel(t, h, "c2", good.URL, 1)
	pushWeightTo(t, h, "c1", 1000)

	w := h.run(t, nil, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
	if w.Code != http.StatusOK {
		t.Fatalf("code=%d body=%s", w.Code, w.Body.String())
	}
	out := w.Body.String()
	if !strings.HasPrefix(out, "data:") || !strings.Contains(out, "真流式") {
		t.Fatalf("应从渠道2 收到 SSE 流：%q", out)
	}
}

// 反向 didNotStream：要整包，渠道回了 event-stream，同样换渠道。
func TestNonStreamReceivingStreamFailsOver(t *testing.T) {
	bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		_, _ = io.WriteString(w, "data: [DONE]\n\n")
	}))
	defer bad.Close()
	good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = io.WriteString(w, `{"id":"ok","choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
	}))
	defer good.Close()

	h := newHarness(t, bad.URL, 0, true, nil)
	addChannel(t, h, "c2", good.URL, 1)
	pushWeightTo(t, h, "c1", 1000)

	w := h.run(t, nil, `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`)
	if w.Code != http.StatusOK || !strings.Contains(w.Body.String(), `"content":"hi"`) {
		t.Fatalf("应从渠道2 收到整包 JSON：code=%d body=%s", w.Code, w.Body.String())
	}
}

// 首 token 计时：流式放行前迟迟不出内容，超时取消本次尝试、换下一个渠道。
func TestFirstTokenTimeoutFailsOver(t *testing.T) {
	bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		w.(http.Flusher).Flush()
		select {
		case <-r.Context().Done(): // 被首 token 计时器取消
		case <-time.After(5 * time.Second):
		}
	}))
	defer bad.Close()
	good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		_, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"及时的内容"}}]}`+"\n\n")
		_, _ = io.WriteString(w, "data: [DONE]\n\n")
	}))
	defer good.Close()

	h := newHarness(t, bad.URL, 0, true, nil)
	h.gw.firstTokenWait = 150 * time.Millisecond
	addChannel(t, h, "c2", good.URL, 1)
	pushWeightTo(t, h, "c1", 1000)

	w := h.run(t, nil, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
	if w.Code != http.StatusOK || !strings.Contains(w.Body.String(), "及时的内容") {
		t.Fatalf("首 token 超时应换渠道：code=%d body=%s", w.Code, w.Body.String())
	}
}

// 失败尝试扫到的 usage 不得计入最终账单：渠道1 的 message_start 式用量随失败被丢弃。
func TestUsageFromFailedTryNotBilled(t *testing.T) {
	bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		// role delta + usage 但没有内容：usage 会被扫进 trySink，随后错误事件导致整次尝试废弃
		_, _ = io.WriteString(w, `data: {"choices":[{"delta":{"role":"assistant"}}],"usage":{"prompt_tokens":500,"completion_tokens":0}}`+"\n\n")
		_, _ = io.WriteString(w, `data: {"error":{"message":"boom"}}`+"\n\n")
	}))
	defer bad.Close()
	good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		_, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"答案"}}]}`+"\n\n")
		_, _ = io.WriteString(w, `data: {"choices":[],"usage":{"prompt_tokens":7,"completion_tokens":3}}`+"\n\n")
		_, _ = io.WriteString(w, "data: [DONE]\n\n")
	}))
	defer good.Close()

	h := newHarness(t, bad.URL, 0, true, nil)
	addChannel(t, h, "c2", good.URL, 1)
	pushWeightTo(t, h, "c1", 1000)

	w := h.run(t, nil, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
	if w.Code != http.StatusOK {
		t.Fatalf("code=%d body=%s", w.Code, w.Body.String())
	}
	h.rec.Close() // usage 异步批量落库，先冲刷再查
	var logs []store.UsageLog
	h.db.Where("status = ?", store.UsageStatusSuccess).Find(&logs)
	if len(logs) != 1 {
		t.Fatalf("成功记账应只有一条，实际 %d", len(logs))
	}
	if logs[0].InputTokens != 7 || logs[0].OutputTokens != 3 {
		t.Fatalf("应只计渠道2 的用量 (7/3)，实际 in=%d out=%d", logs[0].InputTokens, logs[0].OutputTokens)
	}
}

// 响应头已被心跳先行写出后，全部候选失败只能以 SSE 错误事件收尾（状态码已不可改）。
func TestStreamErrorEventWhenHeadAlreadySent(t *testing.T) {
	bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		time.Sleep(200 * time.Millisecond) // 期间心跳先把 200 头写给客户端
		_, _ = io.WriteString(w, `data: {"error":{"message":"boom"}}`+"\n\n")
	}))
	defer bad.Close()

	h := newHarness(t, bad.URL, 0, true, nil)
	h.gw.streamHeartbeat = 40 * time.Millisecond

	w := h.run(t, nil, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
	if w.Code != http.StatusOK {
		t.Fatalf("头已写出，状态应为 200：code=%d body=%s", w.Code, w.Body.String())
	}
	out := w.Body.String()
	if !strings.Contains(out, ": ping") {
		t.Fatalf("扣住期间应有心跳：%q", out)
	}
	if !strings.Contains(out, `"upstream_error"`) {
		t.Fatalf("应以 SSE 错误事件收尾：%q", out)
	}
}

// ---------------------------------------------------------------------------
// 判定函数单测

func TestSSEEventHasContent(t *testing.T) {
	cases := []struct {
		proto string
		line  string
		want  bool
	}{
		// chat：role-only 与空 delta 不算内容，content/reasoning/tool_calls 算
		{"chat", `data: {"choices":[{"delta":{"role":"assistant"}}]}`, false},
		{"chat", `data: {"choices":[{"delta":{}}]}`, false},
		{"chat", `data: {"choices":[{"delta":{"content":"x"}}]}`, true},
		{"chat", `data: {"choices":[{"delta":{"reasoning_content":"想"}}]}`, true},
		{"chat", `data: {"choices":[{"delta":{"tool_calls":[{"id":"1"}]}}]}`, true},
		{"chat", "data: [DONE]", false},
		// messages：message_start / ping 不算，块开启与 delta 算
		{"messages", `data: {"type":"message_start","message":{"usage":{"input_tokens":1}}}`, false},
		{"messages", `data: {"type":"ping"}`, false},
		{"messages", `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`, true},
		{"messages", `data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"t","name":"f"}}`, true},
		{"messages", `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"x"}}`, true},
		{"messages", `data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"想"}}`, true},
		// responses：.delta 与 output_item.added 算，created/completed 不算
		{"responses", `data: {"type":"response.created","response":{}}`, false},
		{"responses", `data: {"type":"response.output_text.delta","delta":"x"}`, true},
		{"responses", `data: {"type":"response.output_item.added","item":{"type":"function_call"}}`, true},
	}
	for i, tc := range cases {
		if got := sseEventHasContent([]byte(tc.line), tc.proto); got != tc.want {
			t.Fatalf("case %d (%s): sseEventHasContent(%s) = %v, want %v", i, tc.proto, tc.line, got, tc.want)
		}
	}
	// 未知协议：任何数据事件都算内容（退化放行，不吞正文）
	if !sseEventHasContent([]byte(`data: {"weird":true}`), "other") {
		t.Fatal("未知协议应放行")
	}
}

func TestSSEUpstreamError(t *testing.T) {
	cases := []struct {
		line    string
		want    bool
		contain string
	}{
		{`data: {"error":{"message":"overloaded","type":"server_error"}}`, true, "overloaded"},
		{`data: {"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}`, true, "Overloaded"},
		{`data: {"type":"error","code":500}`, true, "upstream stream error"},
		{`data: {"type":"response.failed","response":{"error":{"message":"bad tool"}}}`, true, "bad tool"},
		{`data: {"choices":[{"delta":{"content":"hi"}}]}`, false, ""},
		{`data: {"type":"content_block_start","content_block":{"type":"text"}}`, false, ""},
		{`data: {"error":null}`, false, ""},
		{`event: message_start`, false, ""},
	}
	for i, tc := range cases {
		msg, ok := sseUpstreamError([]byte(tc.line))
		if ok != tc.want {
			t.Fatalf("case %d: sseUpstreamError(%s) ok=%v, want %v", i, tc.line, ok, tc.want)
		}
		if ok && !strings.Contains(msg, tc.contain) {
			t.Fatalf("case %d: msg=%q, want contains %q", i, msg, tc.contain)
		}
	}
}

func TestClassifyUpstreamError(t *testing.T) {
	cases := []struct {
		status    int
		body      string
		kind      string
		retryable bool
	}{
		{http.StatusTooManyRequests, `{"error":{"message":"Rate limit reached"}}`, "rate_limited", true},
		{http.StatusTooManyRequests, `{"error":{"message":"You exceeded your current quota, check your plan and billing details"}}`, "quota_exhausted", true},
		{http.StatusTooManyRequests, `{"error":{"code":"1113","message":"余额不足"}}`, "quota_exhausted", true},
		{http.StatusPaymentRequired, `{"error":{"message":"insufficient balance"}}`, "quota_exhausted", true},
		{http.StatusForbidden, `{"error":{"message":"insufficient quota"}}`, "quota_exhausted", true},
		{http.StatusForbidden, `{"error":{"message":"model not allowed"}}`, "server_error", false},
		{http.StatusBadGateway, `{"error":{"message":"upstream gone"}}`, "server_error", true},
		{http.StatusBadRequest, `{"error":{"message":"bad request"}}`, "server_error", false},
	}
	for i, tc := range cases {
		kind, retryable := classifyUpstreamError(tc.status, []byte(tc.body))
		if kind.String() != tc.kind || retryable != tc.retryable {
			t.Fatalf("case %d: classify(%d) = (%s, %v), want (%s, %v)", i, tc.status, kind, retryable, tc.kind, tc.retryable)
		}
	}
}

func TestStreamErrorEventShapes(t *testing.T) {
	chat := string(streamErrorEvent("chat", 502, "boom"))
	if !strings.HasPrefix(chat, "data: ") || !strings.Contains(chat, `"code":"upstream_error"`) || !strings.Contains(chat, "data: [DONE]") {
		t.Fatalf("chat 错误事件形状不符：%q", chat)
	}
	msg := string(streamErrorEvent("messages", 502, "boom"))
	if !strings.HasPrefix(msg, "event: error\ndata: ") {
		t.Fatalf("messages 错误事件形状不符：%q", msg)
	}
	var m map[string]any
	if err := json.Unmarshal([]byte(strings.TrimPrefix(strings.SplitN(msg, "data: ", 2)[1], "\n\n")), &m); err != nil {
		t.Fatalf("messages 错误事件不是合法 JSON: %v (%s)", err, msg)
	}
	resp := string(streamErrorEvent("responses", 502, "boom"))
	if !strings.HasPrefix(resp, "data: ") || strings.Contains(resp, "[DONE]") {
		t.Fatalf("responses 错误事件形状不符：%q", resp)
	}
}