package provider import ( "context" "fmt" "net/http" "net/http/httptest" "sync/atomic" "testing" "modelrouter/internal/config" "modelrouter/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) } }