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)