ModelRouter/internal/provider/retry_policy_test.go

package provider

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"

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

// 上游重试是**内核的通用机制**，判据纯数据（status + body 子串），
// 目的是让"账号级限流、等一会儿就好"这类拒绝不冒到客户端。
//
// 真实事故（2026-10-20）：CodeArts 账户并发会话上限 3 个，连续第 3~4 个请求
// 回 400 TM.00001041，最后一个请求结束后约 44~60 秒才释放；流式更糟 ——
// 上游用 HTTP 200 + error 帧表达拒绝，内核只能报 `empty stream`。
// 参考实现的处理是每 10s 重试一次 chat 直到放行，本测试钉住同一语义。
func TestChatRetriesDeclaredUpstreamRejection(t *testing.T) {
	var hits int32
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		n := atomic.AddInt32(&hits, 1)
		if n <= 2 {
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprint(w, `{"error_code":"TM.00001041","error_msg":"并发会话数已达上限(3个)"}`)
			return
		}
		fmt.Fprint(w, `{"choices":[{"message":{"content":"hi"},"finish_reason":"stop"}]}`)
	}))
	defer up.Close()

	s := src("codearts-like", up.URL, "openai", "m1")
	s.Retry = &config.RetryPolicy{
		OnStatus:       []int{400},
		OnBodyContains: []string{"TM.00001041"},
		Delay:          "1ms",
		MaxAttempts:    4,
	}
	p := newTestProvider(t, s)
	resp, err := p.Chat(context.Background(), &types.ChatRequest{
		Model:    "m1",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("hi")}},
	})
	if err != nil {
		t.Fatalf("命中重试策略后应成功，实得错误: %v", err)
	}
	if resp.Content != "hi" {
		t.Fatalf("content = %q", resp.Content)
	}
	if got := atomic.LoadInt32(&hits); got != 3 {
		t.Fatalf("应恰好尝试 3 次（2 次失败 + 1 次成功），实得 %d", got)
	}
}

// 同一状态码但 body 不匹配 ⇒ **不许重试**（400 也常用于真正的请求错误，
// 无差别重试只会把错误延迟 3 分钟才暴露给用户）。
func TestChatDoesNotRetryUnmatchedBody(t *testing.T) {
	var hits int32
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		atomic.AddInt32(&hits, 1)
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprint(w, `{"error_code":"TM.00001005","error_msg":"无正常可用套餐"}`)
	}))
	defer up.Close()

	s := src("codearts-like", up.URL, "openai", "m1")
	s.Retry = &config.RetryPolicy{
		OnStatus:       []int{400},
		OnBodyContains: []string{"TM.00001041"},
		Delay:          "1ms",
		MaxAttempts:    5,
	}
	p := newTestProvider(t, s)
	_, err := p.Chat(context.Background(), &types.ChatRequest{
		Model:    "m1",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("hi")}},
	})
	if err == nil {
		t.Fatal("不可重试的 body 必须失败")
	}
	if got := atomic.LoadInt32(&hits); got != 1 {
		t.Fatalf("body 不匹配时只能尝试 1 次，实得 %d", got)
	}
}

// 流式：只重试**握手**。命中策略 → 重试直到 200 才开始派发；一旦开始流式
// 就不再重试（重复输出比报错更糟）。
func TestChatStreamRetriesHandshakeOnly(t *testing.T) {
	var hits int32
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		n := atomic.AddInt32(&hits, 1)
		if n == 1 {
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprint(w, `{"error_code":"TM.00001041","error_msg":"并发会话数已达上限(3个)"}`)
			return
		}
		w.Header().Set("Content-Type", "text/event-stream")
		fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"hi\"},\"finish_reason\":\"stop\"}]}\n\n")
		fmt.Fprint(w, "data: [DONE]\n\n")
	}))
	defer up.Close()

	s := src("codearts-like", up.URL, "openai", "m1")
	s.Retry = &config.RetryPolicy{
		OnStatus:       []int{400},
		OnBodyContains: []string{"TM.00001041"},
		Delay:          "1ms",
		MaxAttempts:    4,
	}
	p := newTestProvider(t, s)
	ch, err := p.ChatStream(context.Background(), &types.ChatRequest{
		Model:    "m1",
		Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("hi")}},
	})
	if err != nil {
		t.Fatalf("握手重试后应成功: %v", err)
	}
	var content string
	for ck := range ch {
		content += ck.Content
	}
	if content != "hi" {
		t.Fatalf("content = %q", content)
	}
	if got := atomic.LoadInt32(&hits); got != 2 {
		t.Fatalf("应握手重试 1 次（共 2 次请求），实得 %d", got)
	}
}