diff --git a/internal/config/retry_policy_test.go b/internal/config/retry_policy_test.go new file mode 100644 index 0000000..98ced4c --- /dev/null +++ b/internal/config/retry_policy_test.go @@ -0,0 +1,59 @@ +package config + +import ( + "testing" + "time" +) + +// ShouldRetry 是纯数据判据：状态码列表 + 响应体子串。它决定"这次失败值不值得 +// 在网关内部等一会儿重试"。厂商知识绝不能进内核，所以判据必须能用一个表格测清楚。 +func TestRetryPolicyShouldRetry(t *testing.T) { + cases := []struct { + name string + policy *RetryPolicy + status int + body string + want bool + }{ + {"nil policy", nil, 500, "anything", false}, + {"max_attempts=1 disables", &RetryPolicy{OnStatus: []int{400}, MaxAttempts: 1}, 400, "x", false}, + {"status matches, no body filter", + &RetryPolicy{OnStatus: []int{429}, MaxAttempts: 3}, 429, "", true}, + {"status mismatch", + &RetryPolicy{OnStatus: []int{429}, MaxAttempts: 3}, 400, "429", false}, + {"status + body match", + &RetryPolicy{OnStatus: []int{400}, OnBodyContains: []string{"TM.00001041"}, MaxAttempts: 3}, + 400, `{"error_code":"TM.00001041"}`, true}, + {"status matches but body does not", + &RetryPolicy{OnStatus: []int{400}, OnBodyContains: []string{"TM.00001041"}, MaxAttempts: 3}, + 400, `{"error_code":"TM.00001005"}`, false}, + {"empty substring never matches", + &RetryPolicy{OnStatus: []int{400}, OnBodyContains: []string{""}, MaxAttempts: 3}, 400, "x", false}, + {"any of several body signals", + &RetryPolicy{OnStatus: []int{400, 429}, OnBodyContains: []string{"TM.00001041", "TPM limit"}, MaxAttempts: 3}, + 429, `{"message":"TPM limit exceeded"}`, true}, + } + for _, tc := range cases { + if got := tc.policy.ShouldRetry(tc.status, tc.body); got != tc.want { + t.Errorf("%s: ShouldRetry(%d, %q) = %v, want %v", tc.name, tc.status, tc.body, got, tc.want) + } + } +} + +func TestRetryPolicyDelay(t *testing.T) { + cases := []struct { + name string + policy *RetryPolicy + want time.Duration + }{ + {"nil", nil, 10 * time.Second}, + {"explicit", &RetryPolicy{Delay: "250ms"}, 250 * time.Millisecond}, + {"garbage falls back", &RetryPolicy{Delay: "not-a-duration"}, 10 * time.Second}, + {"negative falls back", &RetryPolicy{Delay: "-5s"}, 10 * time.Second}, + } + for _, tc := range cases { + if got := tc.policy.DelayDuration(); got != tc.want { + t.Errorf("%s: DelayDuration() = %v, want %v", tc.name, got, tc.want) + } + } +} diff --git a/internal/provider/retry_policy_test.go b/internal/provider/retry_policy_test.go new file mode 100644 index 0000000..9326fff --- /dev/null +++ b/internal/provider/retry_policy_test.go @@ -0,0 +1,131 @@ +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) + } +}