Respect provider rate limits and preserve bank connections on throttling

2026-09-10 17:25:09 +02:00
parent 2259db3e85
commit ba3ea6ae5a
14 changed files with 1250 additions and 95 deletions
@@ -0,0 +1,341 @@
+package classification
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"io"
+	"net/http"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+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
+}