mirror of
https://github.com/RuoJi6/ARTEX.git
synced 2026-10-27 17:03:49 +08:00
Merge pull request #13 from RuoJi6/codex/fix-llm-strict-rate-limit
fix: enforce strict LLM HTTP request rate limits
This commit is contained in:
5 files changed
+656
-21
No files matched your search
+21
-6
@@ -250,14 +250,30 @@ func (c Config) Provider() string {
|
||||
return "anthropic"
|
||||
}
|
||||
|
||||
// NewProvider builds an llm.Provider from the config. When a rate is set, the
|
||||
// limiter lives on the single provider instance — so planner + all workers +
|
||||
// main agent (which share this provider) are bounded by one shared rate limit.
|
||||
// NewProvider builds a provider with its own HTTP request limiter. Every HTTP
|
||||
// attempt, including SDK retries, counts toward both rolling rate windows.
|
||||
func (c Config) NewProvider() (llm.Provider, error) {
|
||||
return c.NewProviderWithLimiter(&RequestLimiter{})
|
||||
}
|
||||
|
||||
// NewProviderWithLimiter reuses quota history across provider rebuilds. Callers
|
||||
// sharing a profile must retain the same limiter outside their provider cache.
|
||||
func (c Config) NewProviderWithLimiter(limiter *RequestLimiter) (llm.Provider, error) {
|
||||
client, err := quotaAwareHTTPClient(c.Proxy, c.SessionHeaderKey)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return c.newProvider(client, limiter)
|
||||
}
|
||||
|
||||
func (c Config) newProvider(client *http.Client, limiter *RequestLimiter) (llm.Provider, error) {
|
||||
if limiter == nil {
|
||||
limiter = &RequestLimiter{}
|
||||
}
|
||||
if err := limiter.configure(c.RatePerSecond, c.RatePerMinute); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
client.Transport = rateLimitedTransport{base: client.Transport, limiter: limiter}
|
||||
lc := llm.Config{
|
||||
Format: c.Format,
|
||||
BaseURL: c.BaseURL,
|
||||
@@ -277,9 +293,8 @@ func (c Config) NewProvider() (llm.Provider, error) {
|
||||
lc.RetryInterval = c.Retry.ConnectInterval
|
||||
lc.EmptyResponseRetries = c.Retry.EmptyAttempts
|
||||
lc.EmptyResponseInterval = c.Retry.EmptyInterval
|
||||
if c.RatePerSecond > 0 || c.RatePerMinute > 0 {
|
||||
lc.RateLimit = &llm.RateLimit{PerSecond: c.RatePerSecond, PerMinute: c.RatePerMinute}
|
||||
}
|
||||
// Norma's limiter allows a full-window burst and excludes establishment
|
||||
// retries. Enforce limits on HTTP attempts instead; do not also enable it.
|
||||
return llm.NewProvider(lc)
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,136 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// RequestLimiter counts HTTP attempts in rolling second/minute windows. Keep it
|
||||
// outside provider caches so rebuilding a provider does not reset its quota.
|
||||
// Its zero value is ready for use; limits are set by NewProviderWithLimiter.
|
||||
type RequestLimiter struct {
|
||||
mu sync.Mutex
|
||||
windows [2]requestWindow
|
||||
sent []time.Time
|
||||
changed chan struct{}
|
||||
}
|
||||
|
||||
type requestWindow struct {
|
||||
period time.Duration
|
||||
limit float64
|
||||
}
|
||||
|
||||
func (l *RequestLimiter) configure(perSecond, perMinute float64) error {
|
||||
var windows [2]requestWindow
|
||||
for i, rate := range [2]float64{perSecond, perMinute} {
|
||||
if math.IsNaN(rate) || math.IsInf(rate, 0) {
|
||||
return fmt.Errorf("LLM request rate must be finite")
|
||||
}
|
||||
if rate <= 0 {
|
||||
continue
|
||||
}
|
||||
period := [2]time.Duration{time.Second, time.Minute}[i]
|
||||
// Whole HTTP attempts cannot exceed a fractional window capacity.
|
||||
limit := math.Floor(rate)
|
||||
if rate < 1 {
|
||||
// Preserve fractional rates such as 0.5/s as one attempt per 2s.
|
||||
d := math.Ceil(float64(period) / rate)
|
||||
if d >= float64(math.MaxInt64) {
|
||||
return fmt.Errorf("LLM request rate is too small")
|
||||
}
|
||||
period, limit = time.Duration(d), 1
|
||||
}
|
||||
windows[i] = requestWindow{period: period, limit: limit}
|
||||
}
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
l.windows = windows
|
||||
if l.changed != nil {
|
||||
close(l.changed)
|
||||
}
|
||||
l.changed = make(chan struct{})
|
||||
return nil
|
||||
}
|
||||
|
||||
// wait counts a request only when it may actually proceed. Waiters recheck both
|
||||
// windows under the same lock, so concurrency and cancellation cannot overbook
|
||||
// capacity. Requests in (now-period, now] count toward each rolling window.
|
||||
func (l *RequestLimiter) wait(ctx context.Context) error {
|
||||
if l == nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
for {
|
||||
l.mu.Lock()
|
||||
if err := ctx.Err(); err != nil {
|
||||
l.mu.Unlock()
|
||||
return err
|
||||
}
|
||||
now := time.Now()
|
||||
// Retain at least a minute across limit changes, including while limits
|
||||
// are disabled. Fractional rates may need a longer history.
|
||||
retention := time.Minute
|
||||
for _, w := range l.windows {
|
||||
if w.period > retention {
|
||||
retention = w.period
|
||||
}
|
||||
}
|
||||
first := 0
|
||||
// Keep the last attempt even when older than retention: a subsequent
|
||||
// switch to a fractional rate may need to space requests further apart.
|
||||
for first+1 < len(l.sent) && !l.sent[first].After(now.Add(-retention)) {
|
||||
first++
|
||||
}
|
||||
l.sent = l.sent[first:]
|
||||
var delay time.Duration
|
||||
for _, w := range l.windows {
|
||||
if w.limit == 0 {
|
||||
continue
|
||||
}
|
||||
start := len(l.sent)
|
||||
for start > 0 && l.sent[start-1].After(now.Add(-w.period)) {
|
||||
start--
|
||||
}
|
||||
count := len(l.sent) - start
|
||||
if float64(count) >= w.limit {
|
||||
// If a limit was lowered, enough old requests must expire to
|
||||
// leave room for this attempt, not just the oldest request.
|
||||
d := l.sent[len(l.sent)-int(w.limit)].Add(w.period).Sub(now)
|
||||
if d > delay {
|
||||
delay = d
|
||||
}
|
||||
}
|
||||
}
|
||||
if delay <= 0 {
|
||||
l.sent = append(l.sent, now)
|
||||
l.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
changed := l.changed
|
||||
l.mu.Unlock()
|
||||
timer := time.NewTimer(delay)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
timer.Stop()
|
||||
return ctx.Err()
|
||||
case <-changed:
|
||||
timer.Stop()
|
||||
case <-timer.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type rateLimitedTransport struct {
|
||||
base http.RoundTripper
|
||||
limiter *RequestLimiter
|
||||
}
|
||||
|
||||
func (t rateLimitedTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
if err := t.limiter.wait(req.Context()); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return t.base.RoundTrip(req)
|
||||
}
|
||||
@@ -0,0 +1,446 @@
|
||||
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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
package server
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestProfileRequestLimiterSurvivesCacheInvalidation(t *testing.T) {
|
||||
s := &Server{}
|
||||
limiter := s.requestLimiterForProfile(1)
|
||||
if limiter == s.requestLimiterForProfile(2) || limiter == s.requestLimiterForProfile(0) {
|
||||
t.Fatal("different profiles must have independent quotas")
|
||||
}
|
||||
s.invalidateProfileAgents()
|
||||
if got := s.requestLimiterForProfile(1); got != limiter {
|
||||
t.Fatal("clearing providers reset the profile's rate limiter")
|
||||
}
|
||||
}
|
||||
+38
-15
@@ -103,6 +103,9 @@ type Server struct {
|
||||
provCacheMu sync.Mutex
|
||||
provByProfile map[int64]*provEntry
|
||||
provCacheGen uint64
|
||||
// Quota histories survive provider/agent cache invalidation. Key 0 belongs
|
||||
// to the non-persisted global config; persisted profiles use their DB IDs.
|
||||
requestLimiters map[int64]*agent.RequestLimiter
|
||||
|
||||
// llmHealth is the process-wide circuit-breaker state for LLM failover (轮询).
|
||||
// It deliberately lives OUTSIDE the provider caches: rebuilding the chain
|
||||
@@ -446,7 +449,7 @@ func (s *Server) applyLLM(cfg agent.Config) error {
|
||||
}
|
||||
if prov == nil {
|
||||
var err error
|
||||
prov, err = cfg.NewProvider()
|
||||
prov, err = cfg.NewProviderWithLimiter(s.requestLimiterForProfile(0))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -581,31 +584,51 @@ func (s *Server) providerForProfile(id int64) (llm.Provider, agent.Config, bool)
|
||||
if !ok {
|
||||
return nil, agent.Config{}, false
|
||||
}
|
||||
prov, err := cfg.NewProvider()
|
||||
if err != nil {
|
||||
log.Printf("[engine] build provider for LLM profile %d failed: %v", id, err)
|
||||
return nil, agent.Config{}, false
|
||||
}
|
||||
// Wrap with recorder, tagged with this profile's name.
|
||||
if p, _ := s.m.pg.ProfileByID(id); p != nil {
|
||||
prov = llmrec.Wrap(prov, s.m.PG(), cfg.Model, p.Name, cfg.ThinkingType, cfg.ReasoningEffort, s.m.LLMRecordEnabled)
|
||||
prov = bindSideProvider(prov, cfg, id, p.Name)
|
||||
}
|
||||
profile, _ := s.m.pg.ProfileByID(id)
|
||||
s.provCacheMu.Lock()
|
||||
if generation != s.provCacheGen {
|
||||
s.provCacheMu.Unlock()
|
||||
return s.providerForProfile(id)
|
||||
}
|
||||
if e := s.provByProfile[id]; e != nil { // lost the race → keep the winner
|
||||
prov, cfg = e.prov, e.cfg
|
||||
} else {
|
||||
s.provByProfile[id] = &provEntry{prov: prov, cfg: cfg}
|
||||
log.Printf("[engine] built provider for LLM profile %d (%s / %s)", id, cfg.Provider(), cfg.Model)
|
||||
s.provCacheMu.Unlock()
|
||||
return e.prov, e.cfg, true
|
||||
}
|
||||
// Building the HTTP client does not perform network I/O. Keep this step
|
||||
// under the generation lock so a stale build cannot overwrite the retained
|
||||
// limiter's settings after a newer configuration has already been installed.
|
||||
prov, err := cfg.NewProviderWithLimiter(s.requestLimiterForProfileLocked(id))
|
||||
if err != nil {
|
||||
s.provCacheMu.Unlock()
|
||||
log.Printf("[engine] build provider for LLM profile %d failed: %v", id, err)
|
||||
return nil, agent.Config{}, false
|
||||
}
|
||||
if profile != nil {
|
||||
prov = llmrec.Wrap(prov, s.m.PG(), cfg.Model, profile.Name, cfg.ThinkingType, cfg.ReasoningEffort, s.m.LLMRecordEnabled)
|
||||
prov = bindSideProvider(prov, cfg, id, profile.Name)
|
||||
}
|
||||
s.provByProfile[id] = &provEntry{prov: prov, cfg: cfg}
|
||||
log.Printf("[engine] built provider for LLM profile %d (%s / %s)", id, cfg.Provider(), cfg.Model)
|
||||
s.provCacheMu.Unlock()
|
||||
return prov, cfg, true
|
||||
}
|
||||
|
||||
func (s *Server) requestLimiterForProfile(id int64) *agent.RequestLimiter {
|
||||
s.provCacheMu.Lock()
|
||||
defer s.provCacheMu.Unlock()
|
||||
return s.requestLimiterForProfileLocked(id)
|
||||
}
|
||||
|
||||
func (s *Server) requestLimiterForProfileLocked(id int64) *agent.RequestLimiter {
|
||||
if s.requestLimiters == nil {
|
||||
s.requestLimiters = make(map[int64]*agent.RequestLimiter)
|
||||
}
|
||||
if s.requestLimiters[id] == nil {
|
||||
s.requestLimiters[id] = &agent.RequestLimiter{}
|
||||
}
|
||||
return s.requestLimiters[id]
|
||||
}
|
||||
|
||||
// chatAgentForProfile returns a ChatAgent built from a specific LLM profile, cached
|
||||
// per profile id. Returns nil if the profile is missing or has no API key.
|
||||
func (s *Server) chatAgentForProfile(id int64) *agent.ChatAgent {
|
||||
|
||||
Reference in new issue
Block a user