mirror of
https://github.com/RuoJi6/ARTEX.git
synced 2026-10-27 17:03:49 +08:00
447 lines
13 KiB
Go
447 lines
13 KiB
Go
package agent
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"math"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"testing/synctest"
|
|
"time"
|
|
|
|
"github.com/Autumn-27/norma/llm"
|
|
)
|
|
|
|
func TestRequestLimiterRollingWindows(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
name string
|
|
rps, rpm float64
|
|
capacity int
|
|
window time.Duration
|
|
}{
|
|
{"second", 10, 0, 10, time.Second},
|
|
{"minute_45", 0, 45, 45, time.Minute},
|
|
{"fractional_second", .5, 0, 1, 2 * time.Second},
|
|
{"fractional_minute", 0, .5, 1, 2 * time.Minute},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
synctest.Test(t, func(t *testing.T) {
|
|
l := &RequestLimiter{}
|
|
if err := l.configure(tc.rps, tc.rpm); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
start := time.Now()
|
|
for i := 0; i < tc.capacity; i++ {
|
|
if err := l.wait(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
ctx, cancel := context.WithTimeout(context.Background(), tc.window-time.Nanosecond)
|
|
defer cancel()
|
|
if err := l.wait(ctx); !errors.Is(err, context.DeadlineExceeded) {
|
|
t.Fatalf("extra attempt before window expired: %v", err)
|
|
}
|
|
if err := l.wait(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if elapsed := time.Since(start); elapsed != tc.window {
|
|
t.Fatalf("next attempt at %s, want %s; cancelled wait must not consume quota", elapsed, tc.window)
|
|
}
|
|
})
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestRequestLimiterBothWindowsAndConcurrentCallers(t *testing.T) {
|
|
synctest.Test(t, func(t *testing.T) {
|
|
l := &RequestLimiter{}
|
|
if err := l.configure(3, 5); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var mu sync.Mutex
|
|
var sent []time.Time
|
|
var wg sync.WaitGroup
|
|
start := time.Now()
|
|
for i := 0; i < 12; i++ {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
if err := l.wait(context.Background()); err != nil {
|
|
t.Error(err)
|
|
return
|
|
}
|
|
mu.Lock()
|
|
sent = append(sent, time.Now())
|
|
mu.Unlock()
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
if elapsed := time.Since(start); elapsed != 2*time.Minute {
|
|
t.Fatalf("12 requests under 3/s and 5/min took %s, want 2m", elapsed)
|
|
}
|
|
assertRollingRequestCounts(t, sent, time.Second, 3)
|
|
assertRollingRequestCounts(t, sent, time.Minute, 5)
|
|
})
|
|
}
|
|
|
|
func assertRollingRequestCounts(t *testing.T, sent []time.Time, window time.Duration, limit int) {
|
|
t.Helper()
|
|
for _, end := range sent {
|
|
count := 0
|
|
for _, ts := range sent {
|
|
if ts.After(end.Add(-window)) && !ts.After(end) {
|
|
count++
|
|
}
|
|
}
|
|
if count > limit {
|
|
t.Fatalf("%d attempts in rolling %s, limit %d", count, window, limit)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestRequestLimiterReconfiguration(t *testing.T) {
|
|
synctest.Test(t, func(t *testing.T) {
|
|
l := &RequestLimiter{}
|
|
if err := l.configure(0, 5); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for i := 0; i < 5; i++ {
|
|
if err := l.wait(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
start := time.Now()
|
|
if err := l.configure(0, 2); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
done := make(chan error, 1)
|
|
go func() { done <- l.wait(context.Background()) }()
|
|
synctest.Wait()
|
|
select {
|
|
case err := <-done:
|
|
t.Fatalf("lowering limit reset history: %v", err)
|
|
default:
|
|
}
|
|
// A change must wake waiting requests rather than leave them sleeping
|
|
// until the old minute window expires.
|
|
if err := l.configure(0, 6); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := <-done; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if time.Since(start) != 0 {
|
|
t.Fatal("waiter did not wake on limit increase")
|
|
}
|
|
if err := l.configure(0, 2); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := l.wait(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if time.Since(start) != time.Minute {
|
|
t.Fatal("lower limit did not retain previous attempts")
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestRequestLimiterDisabledAndCancelled(t *testing.T) {
|
|
l := &RequestLimiter{}
|
|
if err := l.configure(0, 0); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for i := 0; i < 100; i++ {
|
|
if err := l.wait(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
if err := l.wait(ctx); !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("got %v", err)
|
|
}
|
|
if len(l.sent) != 100 {
|
|
t.Fatal("cancelled request consumed quota")
|
|
}
|
|
for _, rate := range []float64{math.NaN(), math.Inf(1), math.Inf(-1), math.SmallestNonzeroFloat64} {
|
|
if err := l.configure(rate, 0); err == nil {
|
|
t.Errorf("accepted invalid rate %v", rate)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Exercise the actual Norma retry loops for all formats and both completion
|
|
// modes. In-memory HTTP responses let synctest advance a full minute instantly.
|
|
func TestProviderRateLimitCountsHTTPRetries(t *testing.T) {
|
|
for _, format := range []string{"openai", "anthropic", "openai-responses"} {
|
|
for _, streaming := range []bool{false, true} {
|
|
mode := "complete"
|
|
if streaming {
|
|
mode = "stream"
|
|
}
|
|
t.Run(format+"/"+mode, func(t *testing.T) {
|
|
synctest.Test(t, func(t *testing.T) {
|
|
var sent []time.Time
|
|
client := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
|
|
sent = append(sent, time.Now())
|
|
return &http.Response{StatusCode: 500, Header: make(http.Header),
|
|
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"temporary error"}}`))}, nil
|
|
})}
|
|
cfg := ConfigFrom(format, "test-model", "http://example.invalid", "test-key", "")
|
|
cfg.RatePerSecond, cfg.RatePerMinute = 1, 2
|
|
cfg.Retry.ConnectInterval = time.Millisecond
|
|
prov, err := cfg.newProvider(client, &RequestLimiter{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
start := time.Now()
|
|
req := llm.CompletionRequest{Messages: []llm.Message{llm.UserText("ping")}}
|
|
if streaming {
|
|
for _, streamErr := range prov.Stream(context.Background(), req) {
|
|
err = streamErr
|
|
}
|
|
} else {
|
|
_, _, _, err = prov.Complete(context.Background(), req)
|
|
}
|
|
if err == nil {
|
|
t.Fatal("expected upstream 500")
|
|
}
|
|
if len(sent) != 4 {
|
|
t.Fatalf("HTTP attempts=%d, want initial+3 retries", len(sent))
|
|
}
|
|
assertRollingRequestCounts(t, sent, time.Second, 1)
|
|
assertRollingRequestCounts(t, sent, time.Minute, 2)
|
|
if elapsed := time.Since(start); elapsed != time.Minute+time.Second {
|
|
t.Fatalf("retry attempts took %s, want 1m1s", elapsed)
|
|
}
|
|
})
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestProviderRebuildPreservesRateLimit(t *testing.T) {
|
|
synctest.Test(t, func(t *testing.T) {
|
|
l := &RequestLimiter{}
|
|
var sent []time.Time
|
|
cfg := ConfigFrom("openai", "test-model", "http://example.invalid", "test-key", "")
|
|
cfg.RatePerMinute = 1
|
|
cfg.Retry.ConnectAttempts = -1
|
|
for i := 0; i < 2; i++ {
|
|
client := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
|
|
sent = append(sent, time.Now())
|
|
return &http.Response{StatusCode: 500, Header: make(http.Header), Body: io.NopCloser(strings.NewReader("error"))}, nil
|
|
})}
|
|
prov, err := cfg.newProvider(client, l)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
prov.Complete(context.Background(), llm.CompletionRequest{Messages: []llm.Message{llm.UserText("ping")}})
|
|
}
|
|
if len(sent) != 2 || sent[1].Sub(sent[0]) != time.Minute {
|
|
t.Fatalf("provider rebuild reset quota: %v", sent)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestProviderRateLimitSuccessfulRequests(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
name string
|
|
rps, rpm float64
|
|
calls int
|
|
window time.Duration
|
|
limit int
|
|
}{
|
|
{"minute_45_requests_60", 0, 45, 60, time.Minute, 45},
|
|
{"second_10_requests_15", 10, 0, 15, time.Second, 10},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
synctest.Test(t, func(t *testing.T) {
|
|
var sent []time.Time
|
|
client := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
|
|
sent = append(sent, time.Now())
|
|
return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(
|
|
`{"choices":[{"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}`))}, nil
|
|
})}
|
|
cfg := ConfigFrom("openai", "test-model", "http://example.invalid", "test-key", "")
|
|
cfg.RatePerSecond, cfg.RatePerMinute = tc.rps, tc.rpm
|
|
prov, err := cfg.newProvider(client, &RequestLimiter{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
start := time.Now()
|
|
for i := 0; i < tc.calls; i++ {
|
|
_, _, _, err := prov.Complete(context.Background(), llm.CompletionRequest{Messages: []llm.Message{llm.UserText("ping")}})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
if len(sent) != tc.calls || time.Since(start) != tc.window {
|
|
t.Fatalf("sent %d attempts in %s, want %d in %s", len(sent), time.Since(start), tc.calls, tc.window)
|
|
}
|
|
assertRollingRequestCounts(t, sent, tc.window, tc.limit)
|
|
})
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestProviderRateLimitCountsEmptyResponseRetries(t *testing.T) {
|
|
for _, streaming := range []bool{false, true} {
|
|
name := "complete"
|
|
if streaming {
|
|
name = "stream"
|
|
}
|
|
t.Run(name, func(t *testing.T) {
|
|
synctest.Test(t, func(t *testing.T) {
|
|
var sent []time.Time
|
|
client := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
|
|
sent = append(sent, time.Now())
|
|
body := `{"choices":[{"message":{"role":"assistant","content":""},"finish_reason":"stop"}]}`
|
|
if streaming {
|
|
body = "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"
|
|
}
|
|
return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body))}, nil
|
|
})}
|
|
cfg := ConfigFrom("openai", "test-model", "http://example.invalid", "test-key", "")
|
|
cfg.RatePerMinute = 1
|
|
cfg.Retry.EmptyInterval = time.Millisecond
|
|
prov, err := cfg.newProvider(client, &RequestLimiter{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req := llm.CompletionRequest{Messages: []llm.Message{llm.UserText("ping")}}
|
|
if streaming {
|
|
for range prov.Stream(context.Background(), req) {
|
|
}
|
|
} else {
|
|
prov.Complete(context.Background(), req)
|
|
}
|
|
if len(sent) != 3 {
|
|
t.Fatalf("attempts=%d, want initial + 2 empty retries", len(sent))
|
|
}
|
|
assertRollingRequestCounts(t, sent, time.Minute, 1)
|
|
})
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestRequestLimiterUsesRollingHistory(t *testing.T) {
|
|
synctest.Test(t, func(t *testing.T) {
|
|
l := &RequestLimiter{}
|
|
if err := l.configure(3, 0); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
start := time.Now()
|
|
for i := 0; i < 2; i++ {
|
|
if err := l.wait(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
time.Sleep(750 * time.Millisecond)
|
|
if err := l.wait(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
time.Sleep(350 * time.Millisecond)
|
|
for i := 0; i < 2; i++ {
|
|
if err := l.wait(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
// Only the two t=0 attempts expired. The one at t=.75 still counts.
|
|
if err := l.wait(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if time.Since(start) != 1750*time.Millisecond {
|
|
t.Fatal("window reset lost an unexpired attempt")
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestRateLimitedTransportCancellationDoesNotSend(t *testing.T) {
|
|
synctest.Test(t, func(t *testing.T) {
|
|
l := &RequestLimiter{}
|
|
if err := l.configure(0, 1); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
calls := 0
|
|
tr := rateLimitedTransport{limiter: l, base: roundTripperFunc(func(*http.Request) (*http.Response, error) {
|
|
calls++
|
|
return &http.Response{StatusCode: 200, Body: http.NoBody}, nil
|
|
})}
|
|
req, err := http.NewRequest("POST", "http://example.invalid", nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := tr.RoundTrip(req); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
if _, err := tr.RoundTrip(req.WithContext(ctx)); !errors.Is(err, context.DeadlineExceeded) {
|
|
t.Fatalf("got %v", err)
|
|
}
|
|
if calls != 1 {
|
|
t.Fatalf("cancelled waiter sent HTTP request: %d", calls)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestProviderRateLimitPublicFactory(t *testing.T) {
|
|
for _, format := range []string{"openai", "anthropic", "openai-responses"} {
|
|
for _, streaming := range []bool{false, true} {
|
|
mode := "complete"
|
|
if streaming {
|
|
mode = "stream"
|
|
}
|
|
t.Run(format+"/"+mode, func(t *testing.T) {
|
|
var requests atomic.Int32
|
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
requests.Add(1)
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
io.WriteString(w, `{"error":{"message":"temporary error"}}`)
|
|
}))
|
|
defer upstream.Close()
|
|
cfg := ConfigFrom(format, "test-model", upstream.URL, "test-key", "")
|
|
cfg.RatePerSecond, cfg.RatePerMinute = 1, 1
|
|
cfg.Retry.ConnectAttempts = -1
|
|
limiter := &RequestLimiter{}
|
|
req := llm.CompletionRequest{Messages: []llm.Message{llm.UserText("ping")}}
|
|
for i := 0; i < 2; i++ {
|
|
// A new public provider factory call must keep the old quota.
|
|
prov, err := cfg.NewProviderWithLimiter(limiter)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
ctx := context.Background()
|
|
if i == 1 {
|
|
var cancel context.CancelFunc
|
|
ctx, cancel = context.WithTimeout(ctx, 50*time.Millisecond)
|
|
defer cancel()
|
|
}
|
|
if streaming {
|
|
for _, streamErr := range prov.Stream(ctx, req) {
|
|
err = streamErr
|
|
}
|
|
} else {
|
|
_, _, _, err = prov.Complete(ctx, req)
|
|
}
|
|
if err == nil {
|
|
t.Fatal("expected upstream error or quota wait cancellation")
|
|
}
|
|
if i == 1 && !errors.Is(err, context.DeadlineExceeded) {
|
|
t.Fatalf("got %v, want quota wait deadline", err)
|
|
}
|
|
}
|
|
if got := requests.Load(); got != 1 {
|
|
t.Fatalf("provider rebuild sent %d HTTP requests under 1/min", got)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
}
|