Respect provider rate limits and preserve bank connections on throttling

internal/classification/client.go

@@ -11,18 +11,63 @@ import (
 	"net/http"
 	"net/url"
 	"strings"
+	"sync/atomic"
 	"time"
 	"unicode/utf8"
 
 	"finance-duck/internal/domain"
+	"finance-duck/internal/ratelimit"
 )
 
+// Client configuration must not be mutated concurrently with classification.
+// Do not copy a Client after use; use WithModel to share its rate control safely.
 type Client struct {
 	APIKey        string
 	Model         string
 	IncludeAmount bool
 	HTTPClient    *http.Client
 	BaseURL       string
+
+	rate atomic.Pointer[ratelimit.Controller]
+}
+
+// WithModel snapshots the configuration while sharing the original client's
+// in-flight request gate and provider cooldown, including across model choices.
+func (c *Client) WithModel(model string) *Client {
+	snapshot := &Client{
+		APIKey:        c.APIKey,
+		Model:         model,
+		IncludeAmount: c.IncludeAmount,
+		HTTPClient:    c.HTTPClient,
+		BaseURL:       c.BaseURL,
+	}
+	snapshot.rate.Store(c.rateControl())
+	return snapshot
+}
+
+func (c *Client) rateControl() *ratelimit.Controller {
+	if gate := c.rate.Load(); gate != nil {
+		return gate
+	}
+	gate := &ratelimit.Controller{}
+	if c.rate.CompareAndSwap(nil, gate) {
+		return gate
+	}
+	return c.rate.Load()
+}
+
+// Keep context identity without exposing transport errors containing URLs or
+// response details, including deadlines enforced by http.Client itself.
+func requestContextError(ctx context.Context, err error) error {
+	if cause := ctx.Err(); cause != nil {
+		return cause
+	}
+	for _, cause := range []error{context.Canceled, context.DeadlineExceeded} {
+		if errors.Is(err, cause) {
+			return cause
+		}
+	}
+	return nil
 }
 
 type Proposal struct {
@@ -41,9 +86,12 @@ func (c *Client) Classify(ctx context.Context, facts domain.Facts, data domain.D
 		}
 	}
 	p := Proposal{Enrichment: domain.Fallback(facts)}
+	failError := func(err error) (Proposal, error) {
+		p.Enrichment.Classification = domain.Provenance{Source: "fallback", Timestamp: time.Now().UTC().Format(time.RFC3339), Error: err.Error()}
+		return p, err
+	}
 	fail := func(message string) (Proposal, error) {
-		p.Enrichment.Classification = domain.Provenance{Source: "fallback", Timestamp: time.Now().UTC().Format(time.RFC3339), Error: message}
-		return p, errors.New(message)
+		return failError(errors.New(message))
 	}
 	if _, err := facts.Amount.Minor(); err != nil {
 		return fail("invalid transaction amount")
@@ -64,9 +112,16 @@ func (c *Client) Classify(ctx context.Context, facts domain.Facts, data domain.D
 			return p, nil
 		}
 	}
-	if strings.TrimSpace(c.APIKey) == "" || strings.TrimSpace(c.Model) == "" {
+	apiKey, model := c.APIKey, c.Model
+	includeAmount, baseURL, configuredHTTPClient := c.IncludeAmount, c.BaseURL, c.HTTPClient
+	if strings.TrimSpace(apiKey) == "" || strings.TrimSpace(model) == "" {
 		return fail("AI classification is not configured")
 	}
+	gate := c.rateControl()
+	if err := gate.Acquire(ctx); err != nil {
+		return failError(err)
+	}
+	defer gate.Release()
 	clean := newSanitizer(facts, data, false)
 	merchantClean := newSanitizer(facts, data, true)
 	candidates := retrieve(localDescription, p.Enrichment.Kind, data, clean, merchantClean)
@@ -78,7 +133,7 @@ func (c *Client) Classify(ctx context.Context, facts domain.Facts, data domain.D
 		Amount      *domain.Money `json:"amount,omitempty"`
 		Currency    string        `json:"currency,omitempty"`
 	}{Description: clean(facts.RawDescription), Categories: candidates.categories, Tags: candidates.tags, Merchants: candidates.merchants}
-	if c.IncludeAmount {
+	if includeAmount {
 		prompt.Amount = &facts.Amount
 		// Currency is validated separately rather than copied from arbitrary bank text.
 		if len(facts.Currency) != 3 || strings.IndexFunc(facts.Currency, func(r rune) bool { return r < 'A' || r > 'Z' }) >= 0 {
@@ -91,7 +146,7 @@ func (c *Client) Classify(ctx context.Context, facts domain.Facts, data domain.D
 		return fail("cannot encode classification request")
 	}
 	request := map[string]any{
-		"model":      c.Model,
+		"model":      model,
 		"stream":     false,
 		"max_tokens": 512,
 		// Fail closed: never retry without these controls. No plugins/tools are enabled.
@@ -108,7 +163,7 @@ func (c *Client) Classify(ctx context.Context, facts domain.Facts, data domain.D
 	if err != nil {
 		return fail("cannot encode classification request")
 	}
-	base := strings.TrimRight(c.BaseURL, "/")
+	base := strings.TrimRight(baseURL, "/")
 	if base == "" {
 		base = "https://openrouter.ai/api/v1"
 	}
@@ -119,24 +174,34 @@ func (c *Client) Classify(ctx context.Context, facts domain.Facts, data domain.D
 	if endpoint.Scheme != "https" && !(endpoint.Scheme == "http" && (endpoint.Hostname() == "localhost" || endpoint.Hostname() == "127.0.0.1" || endpoint.Hostname() == "::1")) {
 		return fail("AI endpoint must use HTTPS")
 	}
-	req, err := http.NewRequestWithContext(ctx, http.MethodPost, base+"/chat/completions", bytes.NewReader(body))
-	if err != nil {
-		return fail("cannot create classification request")
-	}
-	req.Header.Set("Authorization", "Bearer "+c.APIKey)
-	req.Header.Set("Content-Type", "application/json")
 	client := http.Client{Timeout: 45 * time.Second}
-	if c.HTTPClient != nil {
-		client = *c.HTTPClient
+	if configuredHTTPClient != nil {
+		client = *configuredHTTPClient
 		if client.Timeout == 0 {
 			client.Timeout = 45 * time.Second
 		}
 	}
 	// Redirects could send sensitive prompts to endpoints with different policies.
 	client.CheckRedirect = func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }
-	resp, err := client.Do(req)
+	resp, err := gate.Do(ctx, func(ctx context.Context) (*http.Response, error) {
+		// Each attempt uses identical serialized bytes, credentials and controls.
+		req, err := http.NewRequestWithContext(ctx, http.MethodPost, base+"/chat/completions", bytes.NewReader(body))
+		if err != nil {
+			return nil, errors.New("cannot create classification request")
+		}
+		req.Header.Set("Authorization", "Bearer "+apiKey)
+		req.Header.Set("Content-Type", "application/json")
+		resp, err := client.Do(req)
+		if err != nil {
+			if cause := requestContextError(ctx, err); cause != nil {
+				return nil, fmt.Errorf("AI request canceled: %w", cause)
+			}
+			return nil, errors.New("AI request failed")
+		}
+		return resp, nil
+	}, true)
 	if err != nil {
-		return fail("AI request failed")
+		return failError(err)
 	}
 	defer resp.Body.Close()
 	if resp.StatusCode != http.StatusOK {
@@ -145,6 +210,9 @@ func (c *Client) Classify(ctx context.Context, facts domain.Facts, data domain.D
 	const maxResponse = 64 * 1024
 	raw, err := io.ReadAll(io.LimitReader(resp.Body, maxResponse+1))
 	if err != nil || len(raw) > maxResponse {
+		if cause := requestContextError(ctx, err); cause != nil {
+			return failError(fmt.Errorf("AI request canceled: %w", cause))
+		}
 		return fail("invalid AI response size")
 	}
 	var envelope struct {
@@ -202,7 +270,7 @@ func (c *Client) Classify(ctx context.Context, facts domain.Facts, data domain.D
 			e.MerchantID = proposed.ID
 		}
 	}
-	e.Classification = domain.Provenance{Source: "openrouter", Model: c.Model, Timestamp: time.Now().UTC().Format(time.RFC3339)}
+	e.Classification = domain.Provenance{Source: "openrouter", Model: model, Timestamp: time.Now().UTC().Format(time.RFC3339)}
 	validationData := data
 	if proposed != nil {
 		validationData.Merchants = append(append([]domain.Merchant{}, data.Merchants...), *proposed)

internal/classification/client_test.go

@@ -290,7 +290,7 @@ func TestUnsafeMerchantProposalRejected(t *testing.T) {
 }
 
 func TestProviderErrorsNeverRelaxPolicyOrEchoResponse(t *testing.T) {
-	for _, status := range []int{302, 400, 401, 404, 429, 500, 503} {
+	for _, status := range []int{302, 400, 401, 402, 403, 404, 500, 503} {
 		t.Run(fmt.Sprint(status), func(t *testing.T) {
 			f, d := fixture()
 			calls := 0

internal/classification/rate_limit_test.go

@@ -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
+}