Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,10 @@ an LRU response cache, and per-visitor + global daily token budgets. See
[docs/ai-plan](docs/ai-plan/index.md) for the phased plan and
[docs/agents.md](docs/agents.md) for the agent quickstart.

AI features require the initiating account owner's explicit consent to the
configured service. Review or revoke it in **Settings → AI privacy**; see
[AI privacy and consent](docs/privacy-ai.md) for data flows and upgrade behavior.

## Install

Prebuilt static Linux binaries (amd64 + arm64, web app and docs embedded) are on the
Expand Down
79 changes: 79 additions & 0 deletions ai/accounting_regression_test.go
Original file line number Diff line number Diff line change
@@ -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)
}
62 changes: 58 additions & 4 deletions ai/ai.go
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,22 @@ func (c *Client) ProviderName() string {
return c.provider.Name()
}

// ProviderBaseURL returns the internal endpoint for consent identity/disclosure.
// Callers must remove credentials, path and query before exposing it to clients.
func (c *Client) ProviderBaseURL() string {
if !c.Enabled() {
return ""
}
switch p := c.provider.(type) {
case *openAIProvider:
return p.baseURL
case *anthropicProvider:
return p.baseURL
default:
return "" // The mock provider has no network endpoint.
}
}

// DefaultModel returns the default model, or the provider's default if unset.
func (c *Client) DefaultModel() string {
if !c.Enabled() {
Expand All @@ -112,6 +128,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)
}
Expand All @@ -120,10 +141,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()
Expand All @@ -133,10 +157,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
}
Expand Down Expand Up @@ -173,12 +206,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.
Expand Down
8 changes: 4 additions & 4 deletions ai/ai_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
103 changes: 85 additions & 18 deletions ai/budget.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package ai

import (
"fmt"
"math"
"sync"
"time"
)
Expand All @@ -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.
Expand All @@ -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 {
Expand All @@ -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()
Expand All @@ -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
}
Expand All @@ -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.
Expand Down Expand Up @@ -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
}
6 changes: 3 additions & 3 deletions ai/digest.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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 ""
}
Expand Down
Loading
Loading