Restore AI budget and authentication fixes that never reached main

ai/accounting_regression_test.go

@@ -0,0 +1,79 @@
+package ai
+
+import (
+	"context"
+	"fmt"
+	"math"
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestClient_BudgetReservesInFlightCalls(t *testing.T) {
+	c := newTestClient(t, &Config{Provider: "mock", VisitorDailyTokenBudget: 10, GlobalDailyTokenBudget: 10})
+	started := make(chan struct{}, 2)
+	release := make(chan struct{})
+	c.Mock().SetHandler(func(req *Request) (*Response, error) {
+		started <- struct{}{}
+		<-release
+		return &Response{Text: "answer", OutputTokens: 10}, nil
+	})
+	first := make(chan error, 1)
+	go func() {
+		_, err := c.Complete(context.Background(), &Request{Feature: FeatureChat, MaxTokens: 10, UserKey: "one"})
+		first <- err
+	}()
+	<-started
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
+	defer cancel()
+	_, err := c.Complete(ctx, &Request{Feature: FeatureDigest, MaxTokens: 10, UserKey: "two"})
+	close(release)
+	require.NoError(t, <-first)
+	require.ErrorIs(t, err, ErrBudgetExceeded)
+	require.Equal(t, int64(10), c.Report("one").Global.Total())
+}
+
+func TestEmbedder_CallsCannotBypassBudget(t *testing.T) {
+	c := newTestClient(t, &Config{Provider: "mock", GlobalDailyTokenBudget: 1})
+	c.budget.Record("user:one", 1, 0)
+	_, err := c.Embedder("model").Embed(context.Background(), []string{"uncached"})
+	require.ErrorIs(t, err, ErrBudgetExceeded)
+	require.Equal(t, 0, c.Mock().EmbedCalls())
+}
+
+func TestBudget_NegativeAndOverflowUsageCannotReopenBudget(t *testing.T) {
+	b := NewBudget(10, 10)
+	b.Record("one", math.MaxInt64, math.MaxInt64)
+	b.Record("one", -100, -100)
+	require.ErrorIs(t, b.Allow("one"), ErrBudgetExceeded)
+	require.GreaterOrEqual(t, b.Report("one").Global.Total(), int64(10))
+}
+
+func TestBudget_ReservationSurvivesMidnightAndFinishesOnce(t *testing.T) {
+	now := time.Date(2026, 10, 13, 23, 59, 59, 0, time.UTC)
+	b := NewBudget(10, 10)
+	b.now = func() time.Time { return now }
+	first, err := b.Reserve("one", 8)
+	require.NoError(t, err)
+	now = now.Add(2 * time.Second)
+	_, err = b.Reserve("two", 3)
+	require.ErrorIs(t, err, ErrBudgetExceeded, "a pending call must not disappear at midnight")
+	first.Finish(3, 2)
+	first.Finish(3, 2)
+	second, err := b.Reserve("two", 5)
+	require.NoError(t, err)
+	second.Finish(1, 0)
+	require.Equal(t, int64(6), b.Report("one").Global.Total())
+}
+
+func TestEmbedder_ChargesProviderReportedUsage(t *testing.T) {
+	c, srv := streamingClient(t, func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprint(w, `{"data":[{"index":0,"embedding":[1,2]}],"usage":{"total_tokens":27}}`)
+	}, time.Second)
+	defer srv.Close()
+	_, err := c.Embedder("model").Embed(context.Background(), []string{"one"}, "user:one")
+	require.NoError(t, err)
+	require.Equal(t, int64(27), c.Report("user:one").Visitor.InputTokens)
+}

ai/ai.go

@@ -112,6 +112,11 @@ func (c *Client) Complete(ctx context.Context, req *Request) (*Response, error)
 	if !c.Enabled() {
 		return nil, ErrDisabled
 	}
+	request := *req
+	req = &request
+	if req.MaxTokens <= 0 {
+		req.MaxTokens = 1024
+	}
 	if req.Model == "" {
 		req.Model = c.modelFor(req.Feature)
 	}
@@ -120,10 +125,13 @@ func (c *Client) Complete(ctx context.Context, req *Request) (*Response, error)
 		metrics.AICacheHits.WithLabelValues(string(req.Feature)).Inc()
 		return cached, nil
 	}
-	if err := c.budget.Allow(req.UserKey); err != nil {
+	reservation, err := c.budget.Reserve(req.UserKey, requestInputTokens(req)+int64(req.MaxTokens))
+	if err != nil {
 		metrics.AIRequestsFailure.WithLabelValues(string(req.Feature)).Inc()
 		return nil, err
 	}
+	input, output := requestInputTokens(req), int64(0)
+	defer func() { reservation.Finish(input, output) }()
 	callCtx, cancel := context.WithTimeout(ctx, c.timeout)
 	defer cancel()
 	start := time.Now()
@@ -133,10 +141,19 @@ func (c *Client) Complete(ctx context.Context, req *Request) (*Response, error)
 		metrics.AIRequestsFailure.WithLabelValues(string(req.Feature)).Inc()
 		return nil, err
 	}
+	if response == nil || response.InputTokens < 0 || response.OutputTokens < 0 {
+		return nil, ErrInvalidResponse
+	}
+	input, output = response.InputTokens, response.OutputTokens
+	if input == 0 && output == 0 {
+		input = requestInputTokens(req)
+		output = int64((len(response.Text) + 3) / 4)
+		response.InputTokens = input
+		response.OutputTokens = output
+	}
 	metrics.AIRequestsSuccess.WithLabelValues(string(req.Feature)).Inc()
 	metrics.AITokensInput.WithLabelValues(string(req.Feature)).Add(float64(response.InputTokens))
 	metrics.AITokensOutput.WithLabelValues(string(req.Feature)).Add(float64(response.OutputTokens))
-	c.budget.Record(req.UserKey, response.InputTokens, response.OutputTokens)
 	c.cache.Store(req, response)
 	return response, nil
 }
@@ -173,12 +190,33 @@ func (c *Client) Mock() *MockProvider {
 // embedTexts computes embeddings via the provider's embeddings endpoint. Only
 // OpenAI-compatible providers implement it; others return ErrStreamingNotSupported-style
 // capability errors which callers map to keyword-only fallbacks.
-func (c *Client) embedTexts(ctx context.Context, model string, texts []string) ([][]float32, error) {
+func (c *Client) embedTexts(ctx context.Context, model string, texts []string, userKey string) ([][]float32, error) {
 	te, ok := c.provider.(TextEmbedder)
 	if !ok {
 		return nil, ErrEmbeddingsNotSupported
 	}
-	return te.EmbedTexts(ctx, model, texts)
+	var bytes int64
+	for _, text := range texts {
+		bytes += int64(len(text))
+	}
+	input := (bytes + 3) / 4
+	reservation, err := c.budget.Reserve(userKey, input)
+	if err != nil {
+		return nil, err
+	}
+	defer func() { reservation.Finish(input, 0) }()
+	callCtx, cancel := context.WithTimeout(ctx, c.timeout)
+	defer cancel()
+	if metered, ok := te.(interface {
+		embedTextsWithUsage(context.Context, string, []string) ([][]float32, int64, error)
+	}); ok {
+		vectors, tokens, err := metered.embedTextsWithUsage(callCtx, model, texts)
+		if tokens > 0 {
+			input = tokens
+		}
+		return vectors, err
+	}
+	return te.EmbedTexts(callCtx, model, texts)
 }
 
 // modelFor resolves the per-feature model override, falling back to the default model.

ai/ai_test.go

@@ -61,15 +61,15 @@ func TestClient_BudgetEnforced(t *testing.T) {
 	mock.SetHandler(func(req *Request) (*Response, error) {
 		return &Response{Text: "ok", InputTokens: 60, OutputTokens: 50}, nil
 	})
-	_, err := client.Complete(context.Background(), &Request{Feature: FeatureChat, Prompt: "1", UserKey: "user:phil"})
+	_, err := client.Complete(context.Background(), &Request{Feature: FeatureChat, Prompt: "1", UserKey: "user:phil", MaxTokens: 1})
 	require.Nil(t, err) // 110 tokens, over visitor budget of 100 — charged after the fact
-	_, err = client.Complete(context.Background(), &Request{Feature: FeatureChat, Prompt: "2", UserKey: "user:phil"})
+	_, err = client.Complete(context.Background(), &Request{Feature: FeatureChat, Prompt: "2", UserKey: "user:phil", MaxTokens: 1})
 	require.ErrorIs(t, err, ErrBudgetExceeded)
 
 	// Other visitors can still spend until the global budget is hit (110/150 used)
-	_, err = client.Complete(context.Background(), &Request{Feature: FeatureChat, Prompt: "3", UserKey: "user:ben"})
+	_, err = client.Complete(context.Background(), &Request{Feature: FeatureChat, Prompt: "3", UserKey: "user:ben", MaxTokens: 1})
 	require.Nil(t, err) // global now at 220
-	_, err = client.Complete(context.Background(), &Request{Feature: FeatureChat, Prompt: "4", UserKey: "user:ben"})
+	_, err = client.Complete(context.Background(), &Request{Feature: FeatureChat, Prompt: "4", UserKey: "user:ben", MaxTokens: 1})
 	require.ErrorIs(t, err, ErrBudgetExceeded)
 
 	report := client.Report("user:phil")

ai/budget.go

@@ -2,6 +2,7 @@ package ai
 
 import (
 	"fmt"
+	"math"
 	"sync"
 	"time"
 )
@@ -13,11 +14,13 @@ type Budget struct {
 	visitorDaily int64 // Max tokens per visitor per day; 0 = unlimited
 	globalDaily  int64 // Max tokens server-wide per day; 0 = unlimited
 
-	mu         sync.Mutex
-	day        string            // UTC day the counters belong to, e.g. "2026-09-24"
-	global     Usage             // Global usage for today
-	perVisitor map[string]*Usage // Per-visitor usage for today
-	now        func() time.Time  // Injectable clock for tests
+	mu              sync.Mutex
+	day             string // UTC day the counters belong to, e.g. "2026-09-24"
+	global          Usage  // Global usage for today
+	reserved        int64
+	visitorReserved map[string]int64
+	perVisitor      map[string]*Usage // Per-visitor usage for today
+	now             func() time.Time  // Injectable clock for tests
 }
 
 // Usage is a token/request counter for one day.
@@ -28,7 +31,7 @@ type Usage struct {
 }
 
 // Total returns input+output tokens.
-func (u *Usage) Total() int64 { return u.InputTokens + u.OutputTokens }
+func (u *Usage) Total() int64 { return saturatingAdd(u.InputTokens, u.OutputTokens) }
 
 // UsageReport is the API-facing view of today's usage.
 type UsageReport struct {
@@ -42,17 +45,18 @@ type UsageReport struct {
 // NewBudget creates a Budget. Budget values of 0 mean unlimited.
 func NewBudget(visitorDaily, globalDaily int64) *Budget {
 	return &Budget{
-		visitorDaily: visitorDaily,
-		globalDaily:  globalDaily,
-		perVisitor:   make(map[string]*Usage),
-		now:          time.Now,
+		visitorDaily:    visitorDaily,
+		globalDaily:     globalDaily,
+		perVisitor:      make(map[string]*Usage),
+		visitorReserved: make(map[string]int64),
+		now:             time.Now,
 	}
 }
 
 // Allow returns ErrBudgetExceeded if the visitor or the global daily budget is already
 // exhausted. It does not reserve tokens; Record charges actual usage afterwards, so a
-// single in-flight request may overshoot its budget. This is deliberate: estimation
-// would either throttle legitimate traffic or be gameable.
+// callers performing provider work use Reserve to include in-flight cost. Actual
+// provider usage can exceed the input-token estimate; subsequent calls are blocked.
 func (b *Budget) Allow(userKey string) error {
 	b.mu.Lock()
 	defer b.mu.Unlock()
@@ -73,9 +77,13 @@ func (b *Budget) Record(userKey string, input, output int64) {
 	b.mu.Lock()
 	defer b.mu.Unlock()
 	b.rollover()
-	b.global.Requests++
-	b.global.InputTokens += input
-	b.global.OutputTokens += output
+	b.record(userKey, input, output)
+}
+
+func (b *Budget) record(userKey string, input, output int64) {
+	b.global.Requests = saturatingAdd(b.global.Requests, 1)
+	b.global.InputTokens = saturatingAdd(b.global.InputTokens, input)
+	b.global.OutputTokens = saturatingAdd(b.global.OutputTokens, output)
 	if userKey == "" {
 		return
 	}
@@ -84,9 +92,9 @@ func (b *Budget) Record(userKey string, input, output int64) {
 		u = &Usage{}
 		b.perVisitor[userKey] = u
 	}
-	u.Requests++
-	u.InputTokens += input
-	u.OutputTokens += output
+	u.Requests = saturatingAdd(u.Requests, 1)
+	u.InputTokens = saturatingAdd(u.InputTokens, input)
+	u.OutputTokens = saturatingAdd(u.OutputTokens, output)
 }
 
 // Report returns today's usage for a visitor key and the configured budgets.
@@ -118,3 +126,62 @@ func (b *Budget) rollover() {
 	b.global = Usage{}
 	b.perVisitor = make(map[string]*Usage)
 }
+
+// Reservation holds an estimated input + maximum output charge while a provider
+// call is in flight. Finish releases it and records actual (or estimated) usage.
+type Reservation struct {
+	budget *Budget
+	key    string
+	tokens int64
+	done   bool
+}
+
+// Reserve checks completed and in-flight usage atomically. It never calls a provider.
+func (b *Budget) Reserve(key string, tokens int64) (*Reservation, error) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	b.rollover()
+	tokens = max(0, tokens)
+	if b.globalDaily > 0 && (b.global.Total() >= b.globalDaily || b.reserved > b.globalDaily-b.global.Total() || tokens > b.globalDaily-b.global.Total()-b.reserved) {
+		return nil, fmt.Errorf("%w: global daily budget", ErrBudgetExceeded)
+	}
+	if key != "" && b.visitorDaily > 0 {
+		used := int64(0)
+		if u := b.perVisitor[key]; u != nil {
+			used = u.Total()
+		}
+		pending := b.visitorReserved[key]
+		if used >= b.visitorDaily || pending > b.visitorDaily-used || tokens > b.visitorDaily-used-pending {
+			return nil, fmt.Errorf("%w: visitor daily budget", ErrBudgetExceeded)
+		}
+	}
+	b.reserved = saturatingAdd(b.reserved, tokens)
+	b.visitorReserved[key] = saturatingAdd(b.visitorReserved[key], tokens)
+	return &Reservation{budget: b, key: key, tokens: tokens}, nil
+}
+
+// Finish releases a reservation once, including when the UTC day rolled over.
+func (r *Reservation) Finish(input, output int64) {
+	b := r.budget
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	if r.done {
+		return
+	}
+	r.done = true
+	b.rollover()
+	b.reserved -= r.tokens
+	b.visitorReserved[r.key] -= r.tokens
+	if b.visitorReserved[r.key] == 0 {
+		delete(b.visitorReserved, r.key)
+	}
+	b.record(r.key, input, output)
+}
+
+func saturatingAdd(a, b int64) int64 {
+	a, b = max(0, a), max(0, b)
+	if a > math.MaxInt64-b {
+		return math.MaxInt64
+	}
+	return a + b
+}

ai/digest.go

@@ -193,7 +193,7 @@ func (d *Digester) Digest(ctx context.Context, userKey string, in *DigestInput)
 	}
 	fmt.Fprintf(&prompt, "):\n\n")
 	if d.embedder != nil && len(messages) >= 3 {
-		prompt.WriteString(buildClusterHints(d.embedder, messages))
+		prompt.WriteString(buildClusterHints(ctx, d.embedder, messages, userKey))
 	}
 	for _, m := range messages {
 		timestamp := time.Unix(m.Time, 0).UTC().Format(time.RFC3339)
@@ -276,12 +276,12 @@ func parseDigest(topic string, messageCount int, text string) (*DigestResult, er
 
 // buildClusterHints computes semantic clusters and renders a short hint block for the
 // prompt. Embedding failure degrades silently to no hints (plain summarization).
-func buildClusterHints(e *Embedder, messages []DigestMessage) string {
+func buildClusterHints(ctx context.Context, e *Embedder, messages []DigestMessage, userKey string) string {
 	texts := make([]string, len(messages))
 	for i, m := range messages {
 		texts[i] = m.Title + "\n" + m.Message
 	}
-	vectors, err := e.Embed(context.Background(), texts)
+	vectors, err := e.Embed(ctx, texts, userKey)
 	if err != nil {
 		return ""
 	}

ai/digest_cluster_test.go

@@ -1,6 +1,7 @@
 package ai
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -34,7 +35,7 @@ func TestBuildClusterHints(t *testing.T) {
 		{ID: "m4", Message: "totally different subject entirely"},
 		{ID: "m5", Message: "another unrelated subject here"},
 	}
-	hints := buildClusterHints(embedder, messages)
+	hints := buildClusterHints(context.Background(), embedder, messages, "")
 	t.Logf("hints: %q", hints)
 	// The deterministic mock vectors may or may not split these into 2+ groups;
 	// the function must never panic and must never fabricate message numbers.
@@ -47,6 +48,6 @@ func TestBuildClusterHints_EmptyOnError(t *testing.T) {
 		return nil, ErrInvalidResponse
 	})
 	embedder := client.Embedder("embed-model")
-	hints := buildClusterHints(embedder, []DigestMessage{{Message: "x"}, {Message: "y"}})
+	hints := buildClusterHints(context.Background(), embedder, []DigestMessage{{Message: "x"}, {Message: "y"}}, "")
 	require.Equal(t, "", hints)
 }