diff --git a/cmd/backend/main.go b/cmd/backend/main.go index 11c9714c5..f25851f24 100644 --- a/cmd/backend/main.go +++ b/cmd/backend/main.go @@ -673,7 +673,7 @@ func main() { inboxTagRepository = repository.NewInboxTagRepository(primaryDB.Pool) if key := config.TypeSafeAPIKey(); key != "" { // One TypeSafe client for every typed judgment in this process. - typeSafeClient = typesafe.NewClient(key) + typeSafeClient = typesafe.NewClient(key, typesafe.WithBillingRedis(cache.Client)) } unsubscribeLinkRepository = repository.NewUnsubscribeLinkRepository(primaryDB.Pool) customDomainRepository = repository.NewCustomDomainRepository(primaryDB.Pool) diff --git a/cmd/consumer/main.go b/cmd/consumer/main.go index af86aaf6b..2efa4372b 100644 --- a/cmd/consumer/main.go +++ b/cmd/consumer/main.go @@ -300,7 +300,7 @@ func main() { // no key is configured, and every feature that reads it stays off. var typeSafeClient *typesafe.Client if key := config.TypeSafeAPIKey(); key != "" { - typeSafeClient = typesafe.NewClient(key) + typeSafeClient = typesafe.NewClient(key, typesafe.WithBillingRedis(redisCache.Client)) } if aiProviderC != nil { replyclassify.SetModelClassifier(func(ctx context.Context, system, user string) (string, error) { diff --git a/docs/content/docs/guides/inbox-tagging.mdx b/docs/content/docs/guides/inbox-tagging.mdx index 134040f19..465c0052a 100644 --- a/docs/content/docs/guides/inbox-tagging.mdx +++ b/docs/content/docs/guides/inbox-tagging.mdx @@ -241,6 +241,14 @@ Only messages stored as a notification by a verdict that never asked the action One call contains every question for a message and currently uses about 1,300 input tokens on the recorded fixture set. A notification recognised offline costs one smaller call, for the action check and any of your questions asked of notifications, and nothing when neither is on. An away message recognised offline costs one small call when its return date is checked, and nothing otherwise. TypeSafe evaluates the questions in parallel, but each question still contributes tokens, so check its current pricing before a large backfill. Output tokens are not billed. +The same configured API key also serves reply classification, campaign-copy checks, generated warmup-content review, form triage, bounce classification and import helpers. These all count against the key's TypeSafe account, not just inbox tagging. Self-hosts use their own key and billing account when configured; there is no shared hosted key supplied by default. + +Operators can filter backend and consumer logs by `event=typesafe_usage`. Each request attempt records its feature, question count and HTTP status (zero when no usable API response was received). Successful responses also record the model and input tokens, with the workspace identifier when available. Sum `input_tokens` on successful responses to compare features; failed attempts have no reported token usage. Logs contain no message content, question text, credentials or API keys. Saved inbox and copy verdicts alone are not an account-wide billing ledger: some features do not persist usage, and the TypeSafe console remains the authority for the balance and charges. + +If TypeSafe returns a billing error (HTTP 402), paid judgments pause for five minutes. The next eligible request after that interval becomes a single recovery probe; a successful probe resumes requests, and a failed probe pauses them again. Backend and consumer coordinate this through their existing Redis connection for the same API key. If Redis is unavailable, each process retains its own bounded pause. This prevents repeated failing requests, but does not replenish credits or lower the price of successful judgments. + +Incoming mail is still saved during the pause. Offline rules and already stored verdicts remain available; unavailable model answers are not invented. After restoring credits, a plain backfill can classify skipped historical mail, and `--recheck-notifications` can check notifications filed without an action answer. Recovery probes happen on demand, not by sending paid synthetic health checks. + A message with no Message-ID is skipped. A database claim prevents two concurrent deliveries of the same Message-ID from making duplicate calls, and an abandoned claim becomes eligible for retry after 15 minutes. ## Related guides diff --git a/internal/app/advanced/service.go b/internal/app/advanced/service.go index 540f582c6..fe830727b 100644 --- a/internal/app/advanced/service.go +++ b/internal/app/advanced/service.go @@ -1901,7 +1901,7 @@ func (s *service) IngestDeliverabilityEvent(ctx context.Context, organizationID if eventType == models.DeliverabilityEventBounce && s.bounceJudge != nil && strings.TrimSpace(req.Reason) != "" && !emailverify.NamesRecipient(req.Reason) { // Bounded: a bounce storm must not queue behind a rate-limited judge. - jctx, cancel := context.WithTimeout(ctx, bounceJudgeTimeout) + jctx, cancel := context.WithTimeout(typesafe.WithUsage(ctx, "bounce_classification", organizationID.String()), bounceJudgeTimeout) v, vErr := bounceclass.Classify(jctx, s.bounceJudge, req.Reason) cancel() switch { diff --git a/internal/app/advisor/detect_test.go b/internal/app/advisor/detect_test.go index 56605b919..3ef6dee0c 100644 --- a/internal/app/advisor/detect_test.go +++ b/internal/app/advisor/detect_test.go @@ -963,6 +963,27 @@ func TestJudgeCopyReadsTheCacheBeforeAsking(t *testing.T) { } } +type unavailableJudgeCache struct{ memoryJudgeCache } + +func (m *unavailableJudgeCache) Put(context.Context, uuid.UUID, string, *copyjudge.Verdict) error { + return errors.New("cache unavailable") +} + +func TestJudgeCopyReusesIdenticalCopyWhenCacheWriteFails(t *testing.T) { + judge := &recordingJudge{} + s := NewService(nil, nil, nil, nil, nil, nil, WithCopyJudge(judge, &unavailableJudgeCache{})).(*service) + snap, _ := judgedSnapshot(nil) + for i := 0; i < 50; i++ { + step := snap.Steps[0] + step.ID = uuid.New() + snap.Steps = append(snap.Steps, step) + } + s.judgeCopy(context.Background(), snap) + if judge.calls != 1 || len(snap.CopyJudgments) != 51 { + t.Fatalf("identical copy cost %d calls and produced %d verdicts", judge.calls, len(snap.CopyJudgments)) + } +} + func TestJudgeCopyCapsFreshJudgmentsAndSurvivesErrors(t *testing.T) { judge := &recordingJudge{} cache := &memoryJudgeCache{} diff --git a/internal/app/advisor/judge_copy.go b/internal/app/advisor/judge_copy.go index c8927f969..b45068019 100644 --- a/internal/app/advisor/judge_copy.go +++ b/internal/app/advisor/judge_copy.go @@ -9,6 +9,7 @@ import ( "github.com/google/uuid" "github.com/warmbly/warmbly/internal/app/copyjudge" + "github.com/warmbly/warmbly/internal/pkg/typesafe" "github.com/warmbly/warmbly/internal/repository" ) @@ -38,17 +39,23 @@ func (s *service) judgeCopy(ctx context.Context, snapshot *repository.AdvisorSna return nil } verdicts := make(map[uuid.UUID]copyjudge.Verdict, len(steps)) + byContent := make(map[string]copyjudge.Verdict, len(steps)) fresh := 0 capped := false for _, sc := range steps { body := copyjudge.Body(sc.step.BodyPlain, sc.step.BodyHTML) hash := copyjudge.ContentHash(sc.step.Subject, body) + if verdict, ok := byContent[hash]; ok { + verdicts[sc.step.ID] = verdict + continue + } cached, err := s.judgeCache.Get(ctx, snapshot.OrganizationID, hash) if err != nil { log.Printf("advisor: copy judgment cache read for org %s: %v", snapshot.OrganizationID, err) } if cached != nil { + byContent[hash] = *cached verdicts[sc.step.ID] = *cached continue } @@ -61,7 +68,7 @@ func (s *service) judgeCopy(ctx context.Context, snapshot *repository.AdvisorSna // Bounded per call, and the first failure ends the run's judging: // Evaluate runs inside the refresh request, and a judge that is down // must not be asked forty times. - jctx, cancel := context.WithTimeout(ctx, copyJudgeTimeout) + jctx, cancel := context.WithTimeout(typesafe.WithUsage(ctx, "copy_judgment", snapshot.OrganizationID.String()), copyJudgeTimeout) v, err := copyjudge.Judge(jctx, s.judge, sc.step.Subject, body) cancel() if err != nil { @@ -75,6 +82,7 @@ func (s *service) judgeCopy(ctx context.Context, snapshot *repository.AdvisorSna log.Printf("advisor: copy judgment cache write for org %s: %v", snapshot.OrganizationID, err) } verdicts[sc.step.ID] = *v + byContent[hash] = *v } if capped { log.Printf("advisor: copy judgment cap of %d reached for org %s; the rest are judged next run", maxCopyJudgmentsPerRun, snapshot.OrganizationID) diff --git a/internal/app/bounceclass/classify.go b/internal/app/bounceclass/classify.go index d8f3267df..9df3261a1 100644 --- a/internal/app/bounceclass/classify.go +++ b/internal/app/bounceclass/classify.go @@ -81,7 +81,7 @@ func Classify(ctx context.Context, asker typesafe.Asker, reason string) (*Verdic questions := map[string]typesafe.Question{ questionCause: typesafe.Choice("Why did the receiving server refuse this email?", causeCriteria), } - resp, err := asker.Ask(ctx, state, questions) + resp, err := asker.Ask(typesafe.WithUsage(ctx, "bounce_classification", ""), state, questions) if err != nil { return nil, err } diff --git a/internal/app/contact/import.go b/internal/app/contact/import.go index d0891f8d4..ef4c28621 100644 --- a/internal/app/contact/import.go +++ b/internal/app/contact/import.go @@ -19,6 +19,7 @@ import ( "github.com/warmbly/warmbly/internal/pkg/emailverify" "github.com/warmbly/warmbly/internal/pkg/listquality" "github.com/warmbly/warmbly/internal/pkg/spreadsheet" + "github.com/warmbly/warmbly/internal/pkg/typesafe" "github.com/warmbly/warmbly/internal/repository" "github.com/warmbly/warmbly/internal/utils" ) @@ -123,7 +124,7 @@ func (s *contactService) SuggestImportMapping(ctx context.Context, orgID uuid.UU if s.columnJudge == nil { return suggested, nil } - return importmap.Infer(ctx, s.columnJudge, suggested, headers, shapes, keys) + return importmap.Infer(typesafe.WithUsage(ctx, "contact_import", orgID.String()), s.columnJudge, suggested, headers, shapes, keys) } // existingCustomFieldKeys is the workspace's custom-field keys for the diff --git a/internal/app/copyjudge/judge.go b/internal/app/copyjudge/judge.go index 59a96bbae..fc73ea0aa 100644 --- a/internal/app/copyjudge/judge.go +++ b/internal/app/copyjudge/judge.go @@ -203,7 +203,7 @@ func Judge(ctx context.Context, asker typesafe.Asker, subject, body string) (*Ve return nil, ErrNoContent } - resp, err := asker.Ask(ctx, st, Questions()) + resp, err := asker.Ask(typesafe.WithUsage(ctx, "copy_judgment", ""), st, Questions()) if err != nil { return nil, err } diff --git a/internal/app/form/service.go b/internal/app/form/service.go index c9082dd48..eaec6e50a 100644 --- a/internal/app/form/service.go +++ b/internal/app/form/service.go @@ -407,7 +407,7 @@ func (s *service) Submit(ctx context.Context, publicID string, answers map[strin // and everything below runs as if the form had it off. if f.TriageEnabled && s.triage != nil { tctx, cancel := context.WithTimeout(ctx, triageTimeout) - if verdict, confidence, ok := triage(tctx, s.triage, f.Fields, data); ok { + if verdict, confidence, ok := triage(typesafe.WithUsage(tctx, "form_triage", f.OrganizationID.String()), s.triage, f.Fields, data); ok { sub.Triage = verdict sub.TriageConfidence = confidence } diff --git a/internal/app/form/triage.go b/internal/app/form/triage.go index 9a2561dc4..b303438a0 100644 --- a/internal/app/form/triage.go +++ b/internal/app/form/triage.go @@ -69,7 +69,7 @@ func triage(ctx context.Context, asker typesafe.Asker, fields []models.FormField if runes < triageMinRunes { return "", 0, false } - resp, err := asker.Ask(ctx, state, triageQuestions) + resp, err := asker.Ask(typesafe.WithUsage(ctx, "form_triage", ""), state, triageQuestions) if err != nil || resp == nil { return "", 0, false } diff --git a/internal/app/importmap/policy.go b/internal/app/importmap/policy.go index f1b592815..1f3c1fc6f 100644 --- a/internal/app/importmap/policy.go +++ b/internal/app/importmap/policy.go @@ -125,7 +125,7 @@ func Infer( ctx, cancel := context.WithTimeout(ctx, Timeout) defer cancel() - resp, err := asker.Ask(ctx, st, questions) + resp, err := asker.Ask(typesafe.WithUsage(ctx, "contact_import", ""), st, questions) if err != nil || resp == nil { return mapping, nil } diff --git a/internal/app/inboxtag/replylayer.go b/internal/app/inboxtag/replylayer.go index f2f4d999c..4dfa28674 100644 --- a/internal/app/inboxtag/replylayer.go +++ b/internal/app/inboxtag/replylayer.go @@ -66,7 +66,7 @@ func ReplyClassifier(asker typesafe.Asker, repo repository.InboxTagRepository) r if !HasContent(state) { return replyclassify.Result{}, false } - resp, err := asker.Ask(ctx, state, map[string]Question{ + resp, err := asker.Ask(typesafe.WithUsage(ctx, "reply_classification", in.OrganizationID.String()), state, map[string]Question{ "sentiment": typesafe.Choice("How does this reply to a cold sales email read?", sentimentCriteria), }) if err != nil { diff --git a/internal/app/inboxtag/service.go b/internal/app/inboxtag/service.go index 1b32f5232..04baf7e5b 100644 --- a/internal/app/inboxtag/service.go +++ b/internal/app/inboxtag/service.go @@ -15,6 +15,7 @@ import ( "github.com/warmbly/warmbly/internal/app/replyclassify" "github.com/warmbly/warmbly/internal/models" "github.com/warmbly/warmbly/internal/pkg/dsn" + "github.com/warmbly/warmbly/internal/pkg/typesafe" "github.com/warmbly/warmbly/internal/repository" ) @@ -192,6 +193,7 @@ func (s *Service) Classify(ctx context.Context, m Message) (Decision, error) { // 3. Check deterministic subject, sender, and any supplied header signals // before spending a model call. + ctx = typesafe.WithUsage(ctx, "inbox_tagging", m.OrganizationID.String()) ws := s.workspace(ctx, m.OrganizationID) facts := Facts{DeterministicKind: DeterministicKind(m, ws.languages)} diff --git a/internal/app/mailboximport/mapping.go b/internal/app/mailboximport/mapping.go index 35e221c15..3b163c96b 100644 --- a/internal/app/mailboximport/mapping.go +++ b/internal/app/mailboximport/mapping.go @@ -330,7 +330,7 @@ func (s *Service) askJev(ctx context.Context, cols []models.MailboxImportColumn, } askCtx, cancel := context.WithTimeout(ctx, 6*time.Second) defer cancel() - resp, err := s.asker.Ask(askCtx, state, questions) + resp, err := s.asker.Ask(typesafe.WithUsage(askCtx, "mailbox_import_mapping", ""), state, questions) if err != nil || resp == nil { return } diff --git a/internal/app/mailboximport/runner.go b/internal/app/mailboximport/runner.go index 72f06222d..66e94315a 100644 --- a/internal/app/mailboximport/runner.go +++ b/internal/app/mailboximport/runner.go @@ -152,6 +152,7 @@ func (s *Service) completeFinished(ctx context.Context) { // process connects, updates or parks one row, and records the outcome. func (s *Service) process(ctx context.Context, w repository.ImportWorkRow) { + ctx = typesafe.WithUsage(ctx, "mailbox_import_errors", w.OrgID.String()) rowCtx, cancel := context.WithTimeout(ctx, time.Duration(config.MailboxImportLeaseSeconds-10)*time.Second) defer cancel() @@ -644,7 +645,7 @@ func (s *Service) refineCause(ctx context.Context, fallback, host, detail string } askCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() - resp, err := s.asker.Ask(askCtx, map[string]string{"server_reply": detail, "mail_host": host}, map[string]typesafe.Question{ + resp, err := s.asker.Ask(typesafe.WithUsage(askCtx, "mailbox_import_errors", ""), map[string]string{"server_reply": detail, "mail_host": host}, map[string]typesafe.Question{ "cause": typesafe.Choice("A mail server refused an app's sign-in with state.server_reply. Which of these is the reason?", criteria), }) out := fallback diff --git a/internal/app/mailboximport/service.go b/internal/app/mailboximport/service.go index 2f28945bc..617cb763d 100644 --- a/internal/app/mailboximport/service.go +++ b/internal/app/mailboximport/service.go @@ -228,7 +228,7 @@ func (s *Service) prepare(ctx context.Context, orgID uuid.UUID, in Input, given } } } - sug := s.suggest(ctx, headers, body, given, saved) + sug := s.suggest(typesafe.WithUsage(ctx, "mailbox_import_mapping", orgID.String()), headers, body, given, saved) emails, domains := emailsAndDomains(body, sug.mapping) existing, xerr := s.mailboxes.FindManyInOrganization(ctx, orgID, emails) diff --git a/internal/app/warmupcontent/judge.go b/internal/app/warmupcontent/judge.go index 5f05c048c..20a7fb265 100644 --- a/internal/app/warmupcontent/judge.go +++ b/internal/app/warmupcontent/judge.go @@ -87,7 +87,7 @@ func judgeThread(ctx context.Context, asker typesafe.Asker, subject, description return nil, fmt.Errorf("warmup judge: complete rendered thread exceeds review limit") } state := boundThreadState(subject, description, messages, judgeStateLimit) - resp, err := asker.Ask(ctx, state, judgeQuestions()) + resp, err := asker.Ask(typesafe.WithUsage(ctx, "warmup_content", ""), state, judgeQuestions()) if err != nil { return nil, err } diff --git a/internal/pkg/typesafe/billing.go b/internal/pkg/typesafe/billing.go new file mode 100644 index 000000000..e090c4711 --- /dev/null +++ b/internal/pkg/typesafe/billing.go @@ -0,0 +1,176 @@ +package typesafe + +import ( + "context" + "crypto/sha256" + "errors" + "fmt" + "net/http" + "sync" + "time" + + "github.com/google/uuid" + "github.com/redis/go-redis/v9" + "github.com/rs/zerolog/log" +) + +const billingCooldown = 5 * time.Minute +const billingProbeLease = 3 * time.Minute +const billingStateTTL = 24 * time.Hour + +var ErrBillingCooldown = errors.New("typesafe: paid judgments paused after a billing failure") + +type BillingCooldownError struct{ RetryAfter time.Duration } + +func (e *BillingCooldownError) Error() string { + return fmt.Sprintf("%s; retry in %s", ErrBillingCooldown, e.RetryAfter.Round(time.Second)) +} + +func (e *BillingCooldownError) Unwrap() error { return ErrBillingCooldown } + +type billingPermit struct { + owner string + shared bool +} + +type billingGate struct { + mu sync.Mutex + until time.Time + owner string + now func() time.Time + shared *redis.Client + key string + lastStoreWarning time.Time +} + +func newBillingGate(apiKey string) *billingGate { + return &billingGate{now: time.Now, key: fmt.Sprintf("typesafe:billing:%x", sha256.Sum256([]byte(apiKey)))} +} + +// WithBillingRedis coordinates backend and consumer recovery for the same key. +func WithBillingRedis(client *redis.Client) Option { + return func(c *Client) { c.billing.shared = client } +} + +var acquireBilling = redis.NewScript(` +local until_ms = tonumber(redis.call('HGET', KEYS[1], 'until')) +if not until_ms then return 0 end +local clock = redis.call('TIME') +local now_ms = clock[1] * 1000 + math.floor(clock[2] / 1000) +if until_ms > now_ms then return until_ms - now_ms end +redis.call('HSET', KEYS[1], 'until', now_ms + ARGV[1], 'owner', ARGV[2]) +redis.call('PEXPIRE', KEYS[1], ARGV[3]) +return -1 +`) + +var finishBilling = redis.NewScript(` +if ARGV[1] == 'block' or redis.call('HGET', KEYS[1], 'owner') == ARGV[2] then + if ARGV[1] == 'recover' then + redis.call('DEL', KEYS[1]) + else + local clock = redis.call('TIME') + local now_ms = clock[1] * 1000 + math.floor(clock[2] / 1000) + redis.call('HSET', KEYS[1], 'until', now_ms + ARGV[3], 'owner', '') + redis.call('PEXPIRE', KEYS[1], ARGV[4]) + end + return 1 +end +return 0 +`) + +func (g *billingGate) acquire(ctx context.Context) (billingPermit, error) { + if err := ctx.Err(); err != nil { + return billingPermit{}, err + } + g.mu.Lock() + remaining := g.until.Sub(g.now()) + g.mu.Unlock() + if remaining > 0 { + return billingPermit{}, &BillingCooldownError{RetryAfter: remaining} + } + if g.shared != nil { + owner := uuid.NewString() + sctx, cancel := context.WithTimeout(ctx, 500*time.Millisecond) + wait, err := acquireBilling.Run(sctx, g.shared, []string{g.key}, billingProbeLease.Milliseconds(), owner, billingStateTTL.Milliseconds()).Int64() + cancel() + if err == nil { + if wait > 0 { + return billingPermit{}, &BillingCooldownError{RetryAfter: time.Duration(wait) * time.Millisecond} + } + if wait == -1 { + return billingPermit{owner: owner, shared: true}, nil + } + return billingPermit{shared: true}, nil + } + g.storeWarning() + } + g.mu.Lock() + defer g.mu.Unlock() + if g.until.IsZero() { + return billingPermit{}, nil + } + now := g.now() + if now.Before(g.until) { + return billingPermit{}, &BillingCooldownError{RetryAfter: g.until.Sub(now)} + } + g.owner = uuid.NewString() + g.until = now.Add(billingProbeLease) + return billingPermit{owner: g.owner}, nil +} + +func (g *billingGate) finish(_ context.Context, p billingPermit, outcome error) { + var apiErr *APIError + billing := asAPIError(outcome, &apiErr) && apiErr.Status == http.StatusPaymentRequired + if !billing && p.owner == "" { + return + } + action := "recover" + if outcome != nil { + action = "defer" + } + if billing { + action = "block" + } + recovered := false + g.mu.Lock() + if billing || (!p.shared && g.owner == p.owner) { + g.owner = "" + if outcome == nil { + g.until = time.Time{} + recovered = true + } else { + g.until = g.now().Add(billingCooldown) + } + } + g.mu.Unlock() + if g.shared != nil { + // A cancelled feature request must still publish its billing result. + ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) + applied, err := finishBilling.Run(ctx, g.shared, []string{g.key}, action, p.owner, billingCooldown.Milliseconds(), billingStateTTL.Milliseconds()).Int64() + cancel() + if err != nil { + g.storeWarning() + } + if p.shared && action == "recover" && (err != nil || applied == 0) { + return + } + if p.shared && action == "recover" { + recovered = true + } + } + if billing { + log.Warn().Str("event", "typesafe_billing_paused").Dur("cooldown", billingCooldown).Msg("TypeSafe paid judgments paused; mail ingestion continues") + } else if recovered { + log.Info().Str("event", "typesafe_billing_recovered").Msg("TypeSafe paid judgments resumed") + } +} + +func (g *billingGate) storeWarning() { + g.mu.Lock() + defer g.mu.Unlock() + if g.now().Sub(g.lastStoreWarning) < time.Minute { + return + } + g.lastStoreWarning = g.now() + log.Warn().Str("event", "typesafe_billing_store_unavailable").Msg("TypeSafe billing cooldown is process-local while Redis is unavailable") +} diff --git a/internal/pkg/typesafe/billing_live_test.go b/internal/pkg/typesafe/billing_live_test.go new file mode 100644 index 000000000..fdc197c5a --- /dev/null +++ b/internal/pkg/typesafe/billing_live_test.go @@ -0,0 +1,130 @@ +package typesafe + +import ( + "context" + "errors" + "net/http" + "os" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" + "github.com/redis/go-redis/v9" +) + +func liveBillingRedis(t *testing.T) *redis.Client { + t.Helper() + url := os.Getenv("TYPESAFE_TEST_REDIS_URL") + if url == "" { + t.Skip("set TYPESAFE_TEST_REDIS_URL for shared billing-gate checks") + } + opts, err := redis.ParseURL(url) + if err != nil { + t.Fatal(err) + } + client := redis.NewClient(opts) + t.Cleanup(func() { _ = client.Close() }) + if err := client.Ping(context.Background()).Err(); err != nil { + t.Fatal(err) + } + return client +} + +func TestLiveSharedBillingGate(t *testing.T) { + r := liveBillingRedis(t) + key := uuid.NewString() + a, b := NewClient(key, WithBillingRedis(r)), NewClient(key, WithBillingRedis(r)) + ctx := context.Background() + t.Cleanup(func() { _ = r.Del(ctx, a.billing.key).Err() }) + a.billing.finish(ctx, billingPermit{shared: true}, &APIError{Status: http.StatusPaymentRequired}) + if _, err := b.billing.acquire(ctx); !errors.Is(err, ErrBillingCooldown) { + t.Fatal("another process ignored the billing pause", err) + } + other := NewClient(uuid.NewString(), WithBillingRedis(r)) + if _, err := other.billing.acquire(ctx); err != nil { + t.Fatal("a different API key was paused", err) + } + if err := r.HSet(ctx, a.billing.key, "until", 1).Err(); err != nil { + t.Fatal(err) + } + var admitted atomic.Int64 + var permits sync.Map + var wg sync.WaitGroup + for i := 0; i < 40; i++ { + wg.Add(1) + go func() { + defer wg.Done() + if p, err := b.billing.acquire(ctx); err == nil { + admitted.Add(1) + permits.Store("probe", p) + } + }() + } + wg.Wait() + if admitted.Load() != 1 { + t.Fatalf("%d shared recovery probes", admitted.Load()) + } + p, _ := permits.Load("probe") + b.billing.finish(ctx, p.(billingPermit), nil) + if exists, err := r.Exists(ctx, a.billing.key).Result(); err != nil || exists != 0 { + t.Fatal("successful probe did not clear the shared pause", err) + } + if _, err := b.billing.acquire(ctx); err != nil { + t.Fatal("recovered key stayed paused", err) + } +} + +func TestLiveSharedBillingLeaseAndCancelledPublication(t *testing.T) { + r := liveBillingRedis(t) + g := newBillingGate(uuid.NewString()) + g.shared = r + ctx := context.Background() + t.Cleanup(func() { _ = r.Del(ctx, g.key).Err() }) + cancelled, cancel := context.WithCancel(ctx) + cancel() + g.finish(cancelled, billingPermit{shared: true}, &APIError{Status: 402}) + if exists, err := r.Exists(ctx, g.key).Result(); err != nil || exists != 1 { + t.Fatal("cancelled request failed to publish billing pause", err) + } + g.mu.Lock() + g.until = time.Time{} + g.mu.Unlock() + _ = r.HSet(ctx, g.key, "until", 1).Err() + old, err := g.acquire(ctx) + if err != nil { + t.Fatal(err) + } + _ = r.HSet(ctx, g.key, "until", 1).Err() + current, err := g.acquire(ctx) + if err != nil { + t.Fatal("expired probe was not replaced", err) + } + g.finish(ctx, old, nil) + if exists, err := r.Exists(ctx, g.key).Result(); err != nil || exists != 1 { + t.Fatal("stale probe cleared a new lease", err) + } + g.finish(ctx, current, errors.New("temporary provider failure")) + if _, err := g.acquire(ctx); !errors.Is(err, ErrBillingCooldown) { + t.Fatal("failed probe did not defer recovery", err) + } + if ttl := r.PTTL(ctx, g.key).Val(); ttl <= 0 || ttl > billingStateTTL { + t.Fatalf("unbounded shared state retention: %v", ttl) + } +} + +func TestLiveBillingRedisUnavailableFallsBackLocally(t *testing.T) { + r := liveBillingRedis(t) + _ = r.Close() + g := newBillingGate("key") + g.shared = r + ctx := context.Background() + if _, err := g.acquire(ctx); err != nil { + t.Fatal("Redis failure disabled healthy classification", err) + } + g.finish(ctx, billingPermit{}, &APIError{Status: 402}) + if _, err := g.acquire(ctx); !errors.Is(err, ErrBillingCooldown) { + t.Fatal("Redis failure lost the local billing pause", err) + } +} diff --git a/internal/pkg/typesafe/billing_test.go b/internal/pkg/typesafe/billing_test.go new file mode 100644 index 000000000..fcd863009 --- /dev/null +++ b/internal/pkg/typesafe/billing_test.go @@ -0,0 +1,143 @@ +package typesafe + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + "time" +) + +const goodAnswer = `{"model":"jev-1.13.0","answers":{"kind":{"type":"choice","choice":"human_reply","confidence":0.9}},"usage":{"input_tokens":1300,"output_tokens":2}}` + +func TestBillingCooldownAndRecovery(t *testing.T) { + var calls atomic.Int64 + var status atomic.Int64 + status.Store(402) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls.Add(1) + if status.Load() != 200 { + w.WriteHeader(int(status.Load())) + _, _ = w.Write([]byte(`{"detail":"billing_error: no credits"}`)) + return + } + _, _ = w.Write([]byte(goodAnswer)) + })) + defer srv.Close() + c := testClient(t, srv) + now := time.Now() + c.billing.now = func() time.Time { return now } + ctx := context.Background() + if _, err := c.Ask(ctx, "x", testQuestions()); err == nil { + t.Fatal("402 accepted") + } + for i := 0; i < 50; i++ { + if resp, err := c.Ask(ctx, "x", testQuestions()); resp != nil || !errors.Is(err, ErrBillingCooldown) { + t.Fatalf("paused call fabricated an answer or lost the pause: %v %v", resp, err) + } + } + if calls.Load() != 1 { + t.Fatal("billing failures retried for each arrival") + } + now = now.Add(billingCooldown) + if _, err := c.Ask(ctx, "x", testQuestions()); err == nil { + t.Fatal("failed probe accepted") + } + if calls.Load() != 2 { + t.Fatal("expired cooldown did not probe") + } + status.Store(200) + now = now.Add(billingCooldown) + if _, err := c.Ask(ctx, "x", testQuestions()); err != nil { + t.Fatal(err) + } + if _, err := c.Ask(ctx, "x", testQuestions()); err != nil { + t.Fatal(err) + } + if calls.Load() != 4 { + t.Fatal("successful probe did not resume normal requests") + } +} + +func TestBillingRecoveryHasOneConcurrentProbe(t *testing.T) { + g := newBillingGate("key") + now := time.Now() + g.now = func() time.Time { return now } + g.finish(context.Background(), billingPermit{}, &APIError{Status: 402}) + now = now.Add(billingCooldown) + var admitted atomic.Int64 + var wg sync.WaitGroup + for i := 0; i < 40; i++ { + wg.Add(1) + go func() { + defer wg.Done() + p, err := g.acquire(context.Background()) + if err == nil && p.owner != "" { + admitted.Add(1) + } + }() + } + wg.Wait() + if admitted.Load() != 1 { + t.Fatalf("%d recovery probes admitted", admitted.Load()) + } +} + +func TestBillingLeaseAndStaleSuccess(t *testing.T) { + g := newBillingGate("key") + now := time.Now() + g.now = func() time.Time { return now } + g.finish(context.Background(), billingPermit{}, &APIError{Status: 402}) + now = now.Add(billingCooldown) + old, err := g.acquire(context.Background()) + if err != nil { + t.Fatal(err) + } + now = now.Add(billingProbeLease) + current, err := g.acquire(context.Background()) + if err != nil { + t.Fatal("abandoned probe did not expire", err) + } + g.finish(context.Background(), old, nil) + if _, err := g.acquire(context.Background()); !errors.Is(err, ErrBillingCooldown) { + t.Fatal("stale probe cleared a newer lease") + } + g.finish(context.Background(), current, errors.New("network unavailable")) + if _, err := g.acquire(context.Background()); !errors.Is(err, ErrBillingCooldown) { + t.Fatal("failed recovery probe released the billing gate") + } +} + +func TestNonBillingFailureDoesNotPause(t *testing.T) { + g := newBillingGate("key") + for _, status := range []int{400, 401, 422, 429, 500, 529} { + g.finish(context.Background(), billingPermit{}, &APIError{Status: status}) + if _, err := g.acquire(context.Background()); err != nil { + t.Fatalf("%d paused billing: %v", status, err) + } + } +} + +func TestBillingRecoveryRetryCanSucceed(t *testing.T) { + var calls int + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls++ + if calls == 1 { + w.WriteHeader(429) + return + } + _, _ = w.Write([]byte(goodAnswer)) + })) + defer srv.Close() + c := testClient(t, srv) + c.billing.until = time.Now().Add(-time.Second) + if _, err := c.Ask(context.Background(), "x", testQuestions()); err != nil { + t.Fatal(err) + } + if !c.billing.until.IsZero() { + t.Fatal("successful retry left billing paused") + } +} diff --git a/internal/pkg/typesafe/client.go b/internal/pkg/typesafe/client.go index 1c85e2014..340352d7a 100644 --- a/internal/pkg/typesafe/client.go +++ b/internal/pkg/typesafe/client.go @@ -11,6 +11,8 @@ import ( "net/http" "strconv" "time" + + "github.com/rs/zerolog/log" ) // Package typesafe is the one client for the TypeSafe API, shared by every @@ -130,17 +132,25 @@ type Client struct { maxAttempts int // sleep is the delay function, swapped in tests so a backoff assertion // does not cost wall-clock seconds. - sleep func(context.Context, time.Duration) error + sleep func(context.Context, time.Duration) error + billing *billingGate } -func NewClient(apiKey string) *Client { - return &Client{ +type Option func(*Client) + +func NewClient(apiKey string, options ...Option) *Client { + c := &Client{ endpoint: defaultEndpoint, apiKey: apiKey, http: &http.Client{Timeout: 30 * time.Second}, maxAttempts: 4, sleep: sleepCtx, + billing: newBillingGate(apiKey), } + for _, option := range options { + option(c) + } + return c } func sleepCtx(ctx context.Context, d time.Duration) error { @@ -248,20 +258,41 @@ func (c *Client) Ask(ctx context.Context, state any, questions map[string]Questi if err != nil { return nil, fmt.Errorf("typesafe: encode request: %w", err) } + permit, err := c.billing.acquire(ctx) + if err != nil { + return nil, err + } + var outcome error + defer func() { c.billing.finish(ctx, permit, outcome) }() var lastErr error for attempt := 0; attempt < c.maxAttempts; attempt++ { if attempt > 0 { wait := backoff(attempt, lastErr) if err := c.sleep(ctx, wait); err != nil { + outcome = err return nil, err } } resp, err := c.once(ctx, payload) + entry := log.Info().Str("event", "typesafe_usage").Str("feature", featureFrom(ctx)).Int("questions", len(questions)).Int("attempt", attempt+1) + if scope, ok := ctx.Value(organizationKey{}).(string); ok && scope != "" { + entry.Str("organization_id", scope) + } if err == nil { + outcome = nil + entry.Int("status", http.StatusOK).Str("model", resp.Model).Int("input_tokens", resp.Usage.InputTokens).Msg("TypeSafe request completed") return resp, nil } + var status *APIError + if asAPIError(err, &status) { + entry.Int("status", status.Status) + } else { + entry.Int("status", 0) + } + entry.Msg("TypeSafe request failed") + outcome = err lastErr = err var apiErr *APIError diff --git a/internal/pkg/typesafe/usage.go b/internal/pkg/typesafe/usage.go new file mode 100644 index 000000000..255ebf328 --- /dev/null +++ b/internal/pkg/typesafe/usage.go @@ -0,0 +1,21 @@ +package typesafe + +import "context" + +type featureKey struct{} +type organizationKey struct{} + +func WithUsage(ctx context.Context, feature, organizationID string) context.Context { + ctx = context.WithValue(ctx, featureKey{}, feature) + if organizationID != "" { + ctx = context.WithValue(ctx, organizationKey{}, organizationID) + } + return ctx +} + +func featureFrom(ctx context.Context) string { + if feature, ok := ctx.Value(featureKey{}).(string); ok && feature != "" { + return feature + } + return "unattributed" +} diff --git a/internal/pkg/typesafe/usage_test.go b/internal/pkg/typesafe/usage_test.go new file mode 100644 index 000000000..843e96e63 --- /dev/null +++ b/internal/pkg/typesafe/usage_test.go @@ -0,0 +1,57 @@ +package typesafe + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/rs/zerolog" + "github.com/rs/zerolog/log" +) + +func TestUsageLogsMetadataWithoutContentOrKey(t *testing.T) { + var output bytes.Buffer + previous := log.Logger + log.Logger = zerolog.New(&output) + t.Cleanup(func() { log.Logger = previous }) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write([]byte(goodAnswer)) + })) + defer srv.Close() + c := testClient(t, srv) + ctx := WithUsage(context.Background(), "inbox_tagging", "workspace-id") + if _, err := c.Ask(ctx, "private message body", testQuestions()); err != nil { + t.Fatal(err) + } + var entry struct { + Event string `json:"event"` + Feature string `json:"feature"` + Organization string `json:"organization_id"` + Tokens int `json:"input_tokens"` + Questions int `json:"questions"` + Status int `json:"status"` + } + if err := json.Unmarshal(output.Bytes(), &entry); err != nil { + t.Fatal(err) + } + if entry.Event != "typesafe_usage" || entry.Feature != "inbox_tagging" || entry.Organization != "workspace-id" || entry.Tokens != 1300 || entry.Questions != 1 || entry.Status != 200 { + t.Fatalf("missing usage attribution: %+v", entry) + } + for _, private := range []string{"test-key", "private message body", "What is this message?"} { + if strings.Contains(output.String(), private) { + t.Fatalf("usage logs leaked %q", private) + } + } +} + +func TestUsageFeaturePreservesWorkspaceScope(t *testing.T) { + ctx := WithUsage(context.Background(), "advisor", "workspace-id") + ctx = WithUsage(ctx, "copy_judgment", "") + if featureFrom(ctx) != "copy_judgment" || ctx.Value(organizationKey{}) != "workspace-id" { + t.Fatal("feature annotation lost workspace attribution") + } +}