397 lines
14 KiB
Go
397 lines
14 KiB
Go
package classification
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"finance-duck/internal/ratelimit"
|
|
)
|
|
|
|
func TestProductionRatePolicyAcrossModelSnapshots(t *testing.T) {
|
|
t.Run("successful classifications stay paced", func(t *testing.T) {
|
|
f, d := fixture()
|
|
arrivals := make(chan time.Time, 2)
|
|
c := mockClient(t, func(w http.ResponseWriter, r *http.Request) {
|
|
arrivals <- time.Now()
|
|
reply(w, validAnswer)
|
|
})
|
|
// Remove only the fixture's fast policy, exercising real lazy production
|
|
// initialization and the same shared controller used by previews.
|
|
c.rate.Store(nil)
|
|
snapshot := c.WithModel("test/preview-model")
|
|
ctx, cancel := context.WithTimeout(context.Background(), 6*time.Second)
|
|
defer cancel()
|
|
for _, client := range []*Client{c, snapshot} {
|
|
p, err := client.Classify(ctx, f, d, true)
|
|
if err != nil || p.Enrichment.Classification.Source != "openrouter" {
|
|
t.Fatalf("production-paced classification failed: %+v, %v", p, err)
|
|
}
|
|
}
|
|
first, second := <-arrivals, <-arrivals
|
|
if gap := second.Sub(first); gap < 3*time.Second-25*time.Millisecond {
|
|
t.Fatalf("model snapshots burst after success: %v", gap)
|
|
}
|
|
})
|
|
t.Run("first failure retains conservative retry deadline", func(t *testing.T) {
|
|
f, d := fixture()
|
|
var calls atomic.Int32
|
|
c := mockClient(t, func(w http.ResponseWriter, r *http.Request) {
|
|
calls.Add(1)
|
|
w.WriteHeader(http.StatusTooManyRequests)
|
|
})
|
|
c.rate.Store(nil)
|
|
snapshot := c.WithModel("test/preview-model")
|
|
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
|
defer cancel()
|
|
_, err := c.Classify(ctx, f, d, true)
|
|
var limit *ratelimit.RateLimitError
|
|
if !errors.Is(err, context.DeadlineExceeded) || !errors.As(err, &limit) {
|
|
t.Fatalf("production retry wait lost cancellation or quota identity: %v", err)
|
|
}
|
|
if delay := time.Until(limit.RetryAt()); delay < 14*time.Second {
|
|
t.Fatalf("production retry used a short per-record fallback: %v", delay)
|
|
}
|
|
probeCtx, probeCancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer probeCancel()
|
|
_, err = snapshot.Classify(probeCtx, f, d, true)
|
|
if !errors.As(err, &limit) || errors.Is(err, context.DeadlineExceeded) || calls.Load() != 1 {
|
|
t.Fatalf("preview discarded production cooldown: %v, calls=%d", err, calls.Load())
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestRateLimitRetryPreservesPrivateRequest(t *testing.T) {
|
|
f, d := fixture()
|
|
f.Counterparty = "Alice Privateperson"
|
|
f.CounterpartyIBAN = "DE89370400440532013000"
|
|
f.RawDescription = "Coffee House Alice Privateperson DE89370400440532013000 private_external -918.27 reference secretpayment"
|
|
var requests [][]byte
|
|
var arrivals []time.Time
|
|
var mu sync.Mutex
|
|
c := mockClient(t, func(w http.ResponseWriter, r *http.Request) {
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
arrivals = append(arrivals, time.Now())
|
|
body, err := io.ReadAll(r.Body)
|
|
if err != nil {
|
|
t.Error(err)
|
|
}
|
|
requests = append(requests, body)
|
|
if r.Method != http.MethodPost || r.URL.Path != "/chat/completions" || r.Header.Get("Authorization") != "Bearer test-secret" || r.Header.Get("Content-Type") != "application/json" {
|
|
t.Error("retry changed authenticated JSON endpoint")
|
|
}
|
|
if len(requests) == 1 {
|
|
w.Header().Set("Retry-After", "2")
|
|
w.WriteHeader(http.StatusTooManyRequests)
|
|
_, _ = io.WriteString(w, "sensitive-provider-response")
|
|
return
|
|
}
|
|
reply(w, validAnswer)
|
|
})
|
|
ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
|
|
defer cancel()
|
|
p, err := c.Classify(ctx, f, d, true)
|
|
if err != nil || p.Enrichment.Classification.Source != "openrouter" || p.Enrichment.Classification.Error != "" {
|
|
t.Fatalf("retry did not recover: %+v, %v", p, err)
|
|
}
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
if len(requests) != 2 || !bytes.Equal(requests[0], requests[1]) {
|
|
t.Fatalf("retry did not reuse identical serialized request: %d attempts", len(requests))
|
|
}
|
|
if gap := arrivals[1].Sub(arrivals[0]); gap < 1950*time.Millisecond {
|
|
t.Fatalf("retried before provider's two-second delay: %v", gap)
|
|
}
|
|
var request struct {
|
|
Provider struct {
|
|
DataCollection string `json:"data_collection"`
|
|
ZDR bool `json:"zdr"`
|
|
Require bool `json:"require_parameters"`
|
|
} `json:"provider"`
|
|
Messages []struct{ Role, Content string } `json:"messages"`
|
|
ResponseFormat struct {
|
|
Type string `json:"type"`
|
|
Schema struct {
|
|
Strict bool `json:"strict"`
|
|
} `json:"json_schema"`
|
|
} `json:"response_format"`
|
|
Plugins json.RawMessage `json:"plugins"`
|
|
}
|
|
if err := json.Unmarshal(requests[0], &request); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if request.Provider.DataCollection != "deny" || !request.Provider.ZDR || !request.Provider.Require || request.ResponseFormat.Type != "json_schema" || !request.ResponseFormat.Schema.Strict || len(request.Plugins) != 0 {
|
|
t.Fatal("retry relaxed private structured routing")
|
|
}
|
|
if len(request.Messages) != 2 {
|
|
t.Fatalf("unexpected message count: %d", len(request.Messages))
|
|
}
|
|
for _, secret := range []string{"alice", "privateperson", "3704", "private_external", "918", "secretpayment", "tx_private", "account_private"} {
|
|
if strings.Contains(strings.ToLower(request.Messages[1].Content), secret) {
|
|
t.Errorf("retried prompt leaked %q", secret)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestRateLimitBackoffExhaustionRetainsSharedCooldown(t *testing.T) {
|
|
f, d := fixture()
|
|
var arrivals []time.Time
|
|
var mu sync.Mutex
|
|
c := mockClient(t, func(w http.ResponseWriter, r *http.Request) {
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
arrivals = append(arrivals, time.Now())
|
|
// An invalid hint must not bypass exponential fallback delays.
|
|
w.Header().Set("Retry-After", "not-a-delay")
|
|
w.WriteHeader(http.StatusTooManyRequests)
|
|
_, _ = io.WriteString(w, "sensitive-provider-response")
|
|
})
|
|
snapshot := c.WithModel("test/preview-model")
|
|
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
|
defer cancel()
|
|
p, err := c.Classify(ctx, f, d, true)
|
|
assertSafeRateLimit(t, err, p.Enrichment.Classification.Error)
|
|
mu.Lock()
|
|
observed := append([]time.Time(nil), arrivals...)
|
|
mu.Unlock()
|
|
if len(observed) != 4 {
|
|
t.Fatalf("expected four bounded attempts, got %d", len(observed))
|
|
}
|
|
for i, delay := range []time.Duration{time.Second, 2 * time.Second, 4 * time.Second} {
|
|
if gap := observed[i+1].Sub(observed[i]); gap < delay-50*time.Millisecond {
|
|
t.Errorf("backoff %d retried too early: %v, need %v", i+1, gap, delay)
|
|
}
|
|
}
|
|
for _, client := range []*Client{snapshot, c} {
|
|
probeCtx, probeCancel := context.WithTimeout(context.Background(), time.Second)
|
|
p, err := client.Classify(probeCtx, f, d, true)
|
|
probeCancel()
|
|
assertSafeRateLimit(t, err, p.Enrichment.Classification.Error)
|
|
mu.Lock()
|
|
calls := len(arrivals)
|
|
mu.Unlock()
|
|
if errors.Is(err, context.DeadlineExceeded) || calls != 4 {
|
|
t.Fatalf("retained cooldown waited or contacted provider: %v, attempts=%d", err, calls)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestRateLimitHTTPDateDoesNotRetryEarly(t *testing.T) {
|
|
f, d := fixture()
|
|
var retryAt time.Time
|
|
var arrivals []time.Time
|
|
var mu sync.Mutex
|
|
c := mockClient(t, func(w http.ResponseWriter, r *http.Request) {
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
arrivals = append(arrivals, time.Now())
|
|
if len(arrivals) == 1 {
|
|
retryAt = time.Now().UTC().Add(3 * time.Second).Truncate(time.Second)
|
|
w.Header().Set("Retry-After", retryAt.Format(http.TimeFormat))
|
|
w.WriteHeader(http.StatusTooManyRequests)
|
|
return
|
|
}
|
|
reply(w, validAnswer)
|
|
})
|
|
ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
|
|
defer cancel()
|
|
if p, err := c.Classify(ctx, f, d, true); err != nil || p.Enrichment.Classification.Source != "openrouter" {
|
|
t.Fatalf("HTTP-date retry failed: %+v, %v", p, err)
|
|
}
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
if len(arrivals) != 2 {
|
|
t.Fatalf("expected one HTTP-date retry, got %d attempts", len(arrivals))
|
|
}
|
|
if arrivals[1].Before(retryAt.Add(-25 * time.Millisecond)) {
|
|
t.Fatalf("retried at %v before HTTP-date %v", arrivals[1], retryAt)
|
|
}
|
|
}
|
|
|
|
func TestRateLimitLongHintsFailFastAndLocalRulesBypassCooldown(t *testing.T) {
|
|
for _, hint := range []string{"600", time.Now().UTC().Add(10 * time.Minute).Format(http.TimeFormat), "9223372036854775807"} {
|
|
t.Run(hint, func(t *testing.T) {
|
|
f, d := fixture()
|
|
var calls atomic.Int32
|
|
c := mockClient(t, func(w http.ResponseWriter, r *http.Request) {
|
|
calls.Add(1)
|
|
w.Header().Set("Retry-After", hint)
|
|
w.WriteHeader(http.StatusTooManyRequests)
|
|
_, _ = io.WriteString(w, "sensitive-provider-response")
|
|
})
|
|
snapshot := c.WithModel("test/preview-model")
|
|
for _, client := range []*Client{c, snapshot} {
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
p, err := client.Classify(ctx, f, d, true)
|
|
cancel()
|
|
assertSafeRateLimit(t, err, p.Enrichment.Classification.Error)
|
|
if errors.Is(err, context.DeadlineExceeded) || calls.Load() != 1 {
|
|
t.Fatalf("long hint waited or permitted an early request: %v, attempts=%d", err, calls.Load())
|
|
}
|
|
}
|
|
d.Merchants[0].UseDefaults = true
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel()
|
|
p, err := snapshot.Classify(ctx, f, d, false)
|
|
if err != nil || p.Enrichment.Classification.Source != "rule" || p.Enrichment.MerchantID != "mer_coffee" || calls.Load() != 1 {
|
|
t.Fatalf("cooldown blocked local rule: %+v, %v, attempts=%d", p, err, calls.Load())
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestRateLimitCancellationClosesBodyAndRetainsCooldown(t *testing.T) {
|
|
f, d := fixture()
|
|
var calls atomic.Int32
|
|
lateRequest := make(chan struct{}, 4)
|
|
c := mockClient(t, func(w http.ResponseWriter, r *http.Request) {
|
|
if calls.Add(1) > 1 {
|
|
lateRequest <- struct{}{}
|
|
}
|
|
w.Header().Set("Retry-After", "1")
|
|
w.WriteHeader(http.StatusTooManyRequests)
|
|
_, _ = io.WriteString(w, "sensitive-provider-response")
|
|
})
|
|
closed := make(chan struct{})
|
|
var closeOnce sync.Once
|
|
transport := c.HTTPClient.Transport
|
|
c.HTTPClient.Transport = rateLimitRoundTripFunc(func(r *http.Request) (*http.Response, error) {
|
|
response, err := transport.RoundTrip(r)
|
|
if err == nil && response.StatusCode == http.StatusTooManyRequests {
|
|
response.Body = &rateLimitNotifyingBody{ReadCloser: response.Body, closed: closed, once: &closeOnce}
|
|
}
|
|
return response, err
|
|
})
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
done := make(chan error, 1)
|
|
go func() {
|
|
_, err := c.Classify(ctx, f, d, true)
|
|
done <- err
|
|
}()
|
|
select {
|
|
case <-closed:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("429 response body was not closed before retry waiting")
|
|
}
|
|
cancel()
|
|
select {
|
|
case err := <-done:
|
|
if !errors.Is(err, context.Canceled) || strings.Contains(err.Error(), "sensitive") {
|
|
t.Fatalf("waiting cancellation did not preserve safe context identity: %v", err)
|
|
}
|
|
case <-time.After(time.Second):
|
|
t.Fatal("retry wait ignored cancellation")
|
|
}
|
|
probeCtx, probeCancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer probeCancel()
|
|
p, err := c.WithModel("test/preview-model").Classify(probeCtx, f, d, true)
|
|
assertSafeRateLimit(t, err, p.Enrichment.Classification.Error)
|
|
if errors.Is(err, context.DeadlineExceeded) || calls.Load() != 1 {
|
|
t.Fatalf("cancelled retry lost cooldown or made a late request: %v, attempts=%d", err, calls.Load())
|
|
}
|
|
select {
|
|
case <-lateRequest:
|
|
t.Fatal("cancelled retry made a request after its timer expired")
|
|
case <-time.After(1200 * time.Millisecond):
|
|
}
|
|
}
|
|
|
|
func TestRateLimitQueuedSnapshotCancellationMakesNoLateRequest(t *testing.T) {
|
|
f, d := fixture()
|
|
entered := make(chan struct{})
|
|
release := make(chan struct{})
|
|
var releaseOnce sync.Once
|
|
unblock := func() { releaseOnce.Do(func() { close(release) }) }
|
|
var calls atomic.Int32
|
|
models := make(chan string, 4)
|
|
c := mockClient(t, func(w http.ResponseWriter, r *http.Request) {
|
|
var request struct{ Model string }
|
|
if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
|
|
t.Error(err)
|
|
}
|
|
models <- request.Model
|
|
if calls.Add(1) == 1 {
|
|
close(entered)
|
|
select {
|
|
case <-release:
|
|
case <-r.Context().Done():
|
|
return
|
|
}
|
|
}
|
|
reply(w, validAnswer)
|
|
})
|
|
defer unblock()
|
|
snapshot := c.WithModel("test/preview-model")
|
|
activeCtx, activeCancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer activeCancel()
|
|
activeDone := make(chan error, 1)
|
|
go func() {
|
|
_, err := c.Classify(activeCtx, f, d, true)
|
|
activeDone <- err
|
|
}()
|
|
select {
|
|
case <-entered:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("first AI request did not start")
|
|
}
|
|
queueCtx, queueCancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
|
_, err := snapshot.Classify(queueCtx, f, d, true)
|
|
queueCancel()
|
|
if !errors.Is(err, context.DeadlineExceeded) || calls.Load() != 1 {
|
|
t.Fatalf("queued snapshot contacted provider or ignored cancellation: %v, attempts=%d", err, calls.Load())
|
|
}
|
|
unblock()
|
|
select {
|
|
case err := <-activeDone:
|
|
if err != nil {
|
|
t.Fatalf("queued cancellation disrupted active request: %v", err)
|
|
}
|
|
case <-time.After(time.Second):
|
|
t.Fatal("active request did not complete")
|
|
}
|
|
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
|
defer cancel()
|
|
p, err := snapshot.Classify(ctx, f, d, true)
|
|
if err != nil || p.Enrichment.Classification.Model != "test/preview-model" || calls.Load() != 2 {
|
|
t.Fatalf("cancelled waiter leaked a request or blocked successor: %+v, %v, attempts=%d", p, err, calls.Load())
|
|
}
|
|
if original, preview := <-models, <-models; original != "test/strict-model" || preview != "test/preview-model" {
|
|
t.Fatalf("serialized requests used wrong models: %q, %q", original, preview)
|
|
}
|
|
}
|
|
|
|
func assertSafeRateLimit(t *testing.T, err error, provenance string) {
|
|
t.Helper()
|
|
if err == nil || !strings.Contains(err.Error(), "429") || provenance == "" || strings.Contains(err.Error(), "sensitive") || strings.Contains(provenance, "sensitive") {
|
|
t.Fatalf("unsafe or missing rate-limit failure: %v, provenance=%q", err, provenance)
|
|
}
|
|
}
|
|
|
|
type rateLimitRoundTripFunc func(*http.Request) (*http.Response, error)
|
|
|
|
func (f rateLimitRoundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) {
|
|
return f(r)
|
|
}
|
|
|
|
type rateLimitNotifyingBody struct {
|
|
io.ReadCloser
|
|
closed chan struct{}
|
|
once *sync.Once
|
|
}
|
|
|
|
func (b *rateLimitNotifyingBody) Close() error {
|
|
err := b.ReadCloser.Close()
|
|
b.once.Do(func() { close(b.closed) })
|
|
return err
|
|
}
|