diff --git a/design/pool-keys.md b/design/pool-keys.md index 50de3a5..42cb4f5 100644 --- a/design/pool-keys.md +++ b/design/pool-keys.md @@ -101,29 +101,54 @@ For the keys probed, both `team_id` and `stripe_customer_id` were returned — we'll snapshot the team id per key so the dedupe is local math, no extra API calls. -### Reset cadence is not surfaced - -The OpenAPI schema for `PlanInfoResponse` and `BillingStatusResponse` -contains no field for reset cadence — no `reset_at`, `next_reset`, -`period`, `cycle`, anything. Whether a key's `total_usage` resets daily, -monthly, or never is purely a property of `payment_plan` and is -invisible from the API. - -Empirically, the partner-tier key tested showed only one bucket -(`2026-05-22`) in `full_history=true` results, consistent with either -"new key first day of use" or "daily-resetting." We can't distinguish -without watching across a UTC midnight rollover. - -**Implication**: don't trust `total_usage` as either lifetime or daily. -Always derive what we need from the timeseries. "Today's spend" comes -from `timeseries(today bucket).total_credits`. "Lifetime usage" comes -from summing all buckets. Whichever interpretation is true for the -underlying plan, our day-by-day arithmetic is correct. - -`remaining_credits` from `/plan-info` is informational only — we display -it but don't gate on it, because if the plan resets daily, "remaining" -is for the current period and resets at midnight on pioneer's clock, -which we don't see directly. +### Reset cadence is daily, confirmed + +The OpenAPI schema for `PlanInfoResponse` contains no `reset_at` field, +but empirical data from a `pro`-plan key across three days resolves the +ambiguity: + +``` +2026-05-20 68010 credits used ($680) +2026-05-21 33815 credits used ($338) +2026-05-22 100007 credits used ($1000, ~$0.07 over limit) +total_usage = 0 on the morning of 2026-05-23 (just renewed) +``` + +`credit_limit` is a **daily cap that resets at UTC midnight**, not a +lifetime allowance. `total_usage` and `remaining_credits` from +`/plan-info` reflect the current day only. + +Pioneer allows a small overshoot before cutting off a key — the pro-plan +key hit 100007 credits against a 100000 limit before going 401. Don't +assume the gate is hard; budget 1-2% headroom when computing "will this +key cover the next request." + +**Known plans so far:** + +| plan | credit_limit | daily cap USD | has_payment_method | +|---|---|---|---| +| `partner` | 40000 | $400 | false | +| `pro` | 100000 | $1000 | true | + +**We only accept `pro` plan keys.** `partner` and any unknown plan are +rejected at add time with a clear error ("only pro-plan pioneer keys are +supported"). This keeps the pool homogeneous and avoids surprises from +plans we haven't characterized. + +`total_usage` is safe to use as "spent today." `remaining_credits` is +safe to use as "remaining today." Both reset at UTC midnight. + +**$10 buffer rule.** The picker treats a key as exhausted when +`remaining_credits < 1000` (1000 pioneer credits = $10). This gives +headroom for in-flight requests to settle and for the ~$0.07 overshoot +pioneer tolerates before cutting the key. The constant is +`PoolKeyBufferCredits = 1000` in `internal/pool/pool.go`, which +corresponds to `10_000_000` potluck micros. + +Implication for the reconciler: `total_usage` from `/plan-info` is +exactly "spent today" — cheaper than a separate timeseries call for +that specific number. We still pull the timeseries for historical data +(charts, per-day attribution across multiple days). ### Per-request join is fuzzy @@ -228,25 +253,43 @@ That's not useful — we always send a key. So when we get a 401 mid-flight or during the reconciler probe, we genuinely cannot tell whether the key is dead or just napping until the next reset. -**Implications:** - -- Adding a real-but-tapped-out key to the pool currently fails with - "rejected, double-check it." It would work tomorrow. The "add key" - flow needs a soft acceptance path: store the key as - `pending_validation`, retry on the reconciler's schedule, activate - when it works. -- A 401 during normal operation must not be fatal. Mark the key - `pioneer_health = 'unauthorized'`, record the timestamp, exclude - from picker, retry tomorrow at midnight UTC + jitter. Reactivate - on first success. -- After N consecutive days of 401 (default: 14), give up and mark - the key `revoked` for real. By that point either the key is dead - or the user has stopped paying their pioneer bill, and we shouldn't - keep probing forever. -- The picker needs a circuit breaker: if Pick() hands out a key and - the actual chat call returns 401, mark the key unhealthy *during - the request*, immediately Pick() again, and retry once. Only fail - the user request if the entire pool is exhausted. +### 503 means "auth service down" + +Pioneer also returns `503 Service Unavailable` with a distinct body: + +```json +{"detail": "Authentication service temporarily unavailable. Please retry shortly."} +``` + +This is not a key problem — it's pioneer being down. It must be handled +differently from 401: don't mark the key unhealthy, don't count it +toward the consecutive-failure counter, just skip and retry on the next +tick. Log it at WARN level so we notice if it's sustained. + +**Three distinct failure modes on any pioneer call:** + +| Status | Body pattern | Meaning | Our action | +|---|---|---|---| +| `401` | `"Invalid API key..."` | Key invalid or exhausted | Soft-mark unhealthy, retry tomorrow | +| `401` | `{"message":"Authentication required"...}` | No auth sent | Bug in our code | +| `503` | `"Authentication service temporarily..."` | Pioneer auth down | Skip this tick, retry next | +| `5xx` other | varies | Pioneer backend down | Skip, retry next; don't mark key | + +**Implications for the reconciler and picker:** + +- 401 → `pioneer_health = 'unauthorized'`, record `pioneer_unhealthy_since`, exclude from picker, retry tomorrow +- 503 / other 5xx → leave health state unchanged, retry next scheduled tick +- Adding a real-but-tapped-out key currently fails with "rejected, + double-check it." It would work tomorrow. The "add key" flow needs a + soft acceptance path: store the key as `pending_validation`, retry on + the reconciler's schedule, activate when it works. +- The picker needs a circuit breaker: if Pick() hands out a key and the + actual chat call returns 401, mark the key unhealthy *during the + request*, immediately Pick() again, retry once. Only fail the user + request if the entire pool is exhausted. +- After N consecutive days of 401 (default: 14), give up and mark the + key `revoked` for real. By that point either the key is dead or the + user has stopped paying their pioneer bill. This is the single most important behavioral change in the new design — without it, we permanently lose contributed keys to transient limits. @@ -352,8 +395,8 @@ Background goroutine, 10-minute ticker. Per active key: - Otherwise: attribute to the key owner ("off-platform usage"). 5. Update `user_daily_spend` rows in a transaction. 6. Set `pool_keys.last_billing_sync_at = now()`. -7. Update `today_micros` cache from - `GET /billing/usage/timeseries?full_history=false&interval_minutes=1440`. +7. Update `today_micros` cache directly from `total_usage` in the + `/plan-info` response (already fetched in step 1 — no extra call). On any 401 during steps 1-7: mark the key `pioneer_health = 'unauthorized'`, set `pioneer_unhealthy_since = now()` if not already @@ -708,11 +751,10 @@ real data. - **Should the recompute button be rate-limited?** Pressing it in a tight loop is harmless (idempotent, no API calls), but visually noisy. Cooldown on the button is probably enough. -- **What about pioneer's per-team budget reset cycle?** If we discover - some keys reset weekly or monthly, we might want to surface that on - the key card. For now we assume daily, which is the most generous - reading and won't cause surprises (worst case "spent today" looks - smaller than reality and the cap is over-conservative). +- **What about non-pro plan keys added in the future?** If pioneer + introduces new plans worth supporting, update `probePioneerBilling`'s + allowlist and add a row to the plans table above. Don't silently + accept unknown plans — fail loud so we can characterize them first. - **Failed pioneer requests** — pioneer returns 5xx on overload; today we just propagate to the user. We may want to retry with a different pool key on 5xx the same way we do on 401. diff --git a/server/cmd/server/main.go b/server/cmd/server/main.go index 8775457..f29c0fd 100644 --- a/server/cmd/server/main.go +++ b/server/cmd/server/main.go @@ -106,6 +106,12 @@ func main() { log.Fatal("pool: init failed", "err", err) } + // Start the background reconciler: syncs pioneer billing data every 10 min. + reconcilerCtx, reconcilerCancel := context.WithCancel(context.Background()) + defer reconcilerCancel() + reconciler := pool.NewReconciler(q, keyPool.Decrypt, log.Default()) + go reconciler.Run(reconcilerCtx) + // Hack Club Auth client. Nil when unconfigured; the handlers degrade // gracefully and return 503 with a friendly note. var hcaClient *hca.Client @@ -136,12 +142,11 @@ func main() { r.Post("/auth/logout", hcaLogoutHandler(authSvc, cfg.IsProduction())) apiSrv := &web.Server{ - Q: q, - Auth: authSvc, - Ledger: ledg, - Hub: hub, - Pool: keyPool, - PioneerAPIKey: cfg.Pioneer.APIKey, + Q: q, + Auth: authSvc, + Ledger: ledg, + Hub: hub, + Pool: keyPool, } v1Srv := &v1.Server{ Q: q, diff --git a/server/db/migrations/00006_pool_keys_v2.sql b/server/db/migrations/00006_pool_keys_v2.sql new file mode 100644 index 0000000..34f9d35 --- /dev/null +++ b/server/db/migrations/00006_pool_keys_v2.sql @@ -0,0 +1,140 @@ +-- +goose Up +-- +goose StatementBegin + +-- Extend pool_keys with two-budget model, pioneer health tracking, and +-- pending-validation support. All changes are additive; existing rows +-- and queries keep working unchanged until the follow-up migration drops +-- daily_limit_micros. +-- +-- Two-budget model: +-- max_micros = owner's daily ceiling (replaces daily_limit_micros) +-- shared_micros = portion donated to the shared pool (≤ max_micros) +-- private = max_micros - shared_micros (reserved for owner only) +-- +-- Pioneer health: +-- pioneer_health = 'unknown' | 'healthy' | 'unauthorized' +-- A 401 from pioneer means exhausted OR revoked — we can't tell which. +-- Mark unauthorized, retry daily. After 14 days consecutive, set revoked_at. +-- A 503 from pioneer means their auth service is down — don't touch health. +-- +-- Only pro-plan keys are accepted (credit_limit = 100,000 credits = $1000/day). + +ALTER TABLE pool_keys ADD COLUMN max_micros INTEGER NOT NULL DEFAULT 1000000000; +ALTER TABLE pool_keys ADD COLUMN shared_micros INTEGER NOT NULL DEFAULT 1000000000; +ALTER TABLE pool_keys ADD COLUMN pioneer_team_id TEXT; +ALTER TABLE pool_keys ADD COLUMN pioneer_payment_plan TEXT; +ALTER TABLE pool_keys ADD COLUMN pioneer_credit_limit_micros INTEGER; +ALTER TABLE pool_keys ADD COLUMN pioneer_remaining_micros INTEGER; +ALTER TABLE pool_keys ADD COLUMN pioneer_health INTEGER NOT NULL DEFAULT 0; + -- 0=unknown 1=healthy 2=unauthorized +ALTER TABLE pool_keys ADD COLUMN pioneer_unhealthy_since INTEGER; +ALTER TABLE pool_keys ADD COLUMN pending_validation INTEGER NOT NULL DEFAULT 0; +ALTER TABLE pool_keys ADD COLUMN last_billing_sync_at INTEGER; +ALTER TABLE pool_keys ADD COLUMN revoked_at INTEGER; + +-- Backfill: seed max/shared from the old daily_limit_micros. +UPDATE pool_keys SET + max_micros = daily_limit_micros, + shared_micros = daily_limit_micros; + +-- potluck_requests: one row per chat completion we proxy, written at +-- request start and updated at finish. The reconciler matches pioneer +-- billing rows against this table to attribute cost to the right user. +CREATE TABLE potluck_requests ( + id TEXT PRIMARY KEY, -- our uuid + user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE, + api_key_id TEXT REFERENCES api_keys(id) ON DELETE SET NULL, + pool_key_id TEXT REFERENCES pool_keys(id) ON DELETE SET NULL, + surface TEXT NOT NULL, -- 'web' | 'v1' + model TEXT NOT NULL, + started_at INTEGER NOT NULL, -- unix seconds + finished_at INTEGER, -- NULL while in flight + prompt_tokens INTEGER, + completion_tokens INTEGER, + total_tokens INTEGER, + status TEXT NOT NULL DEFAULT 'pending', -- 'pending' | 'done' | 'error' | 'canceled' + error_code TEXT +) STRICT; +CREATE INDEX potluck_requests_by_user ON potluck_requests(user_id, started_at DESC); +CREATE INDEX potluck_requests_by_key ON potluck_requests(pool_key_id, finished_at); + +-- pool_key_billing_rows: one row per pioneer billing log entry ingested. +-- Idempotent on pioneer's id. Attribution tracks which potluck user +-- caused the spend (NULL = key owner, for off-platform or judge fallback). +CREATE TABLE pool_key_billing_rows ( + id TEXT PRIMARY KEY, -- pioneer's billing row UUID + pool_key_id TEXT NOT NULL REFERENCES pool_keys(id) ON DELETE CASCADE, + pioneer_created_at INTEGER NOT NULL, -- unix seconds + credit_micros INTEGER NOT NULL, -- credit_usage * 10000 + cost_micros INTEGER NOT NULL, -- cost_usd * 1000000 + token_usage INTEGER NOT NULL, + model TEXT NOT NULL, + endpoint TEXT NOT NULL, + attributed_user_id TEXT REFERENCES users(id), -- NULL = key owner + -- attribution: 0=matched 1=judge_paired 2=owner_fallback 3=duplicate + attribution INTEGER NOT NULL, + is_duplicate INTEGER NOT NULL DEFAULT 0, -- 1 when attribution=3 + matched_request_id TEXT REFERENCES potluck_requests(id), + ingested_at INTEGER NOT NULL +) STRICT; +CREATE INDEX pool_key_billing_by_key ON pool_key_billing_rows(pool_key_id, pioneer_created_at DESC); +CREATE INDEX pool_key_billing_by_user ON pool_key_billing_rows(attributed_user_id, pioneer_created_at DESC); + +-- user_daily_spend: per-user-per-UTC-day spend split into shared vs +-- private. Updated by the reconciler in the same transaction as +-- pool_key_billing_rows inserts. +-- day = unix_seconds / 86400. +CREATE TABLE user_daily_spend ( + user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE, + day INTEGER NOT NULL, + shared_spent_micros INTEGER NOT NULL DEFAULT 0, + private_spent_micros INTEGER NOT NULL DEFAULT 0, + PRIMARY KEY (user_id, day) +) STRICT, WITHOUT ROWID; + +-- user_daily_allowances: how much of the shared pool each user may +-- spend today. Set by the recompute button. Allowances only ever grow +-- on a recompute (no claw-back). If no row exists for today, the gate +-- falls back to a fair-share calculation on the fly. +CREATE TABLE user_daily_allowances ( + user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE, + day INTEGER NOT NULL, + shared_allowance_micros INTEGER NOT NULL, + set_at INTEGER NOT NULL, + set_by_user_id TEXT NOT NULL REFERENCES users(id), + PRIMARY KEY (user_id, day) +) STRICT, WITHOUT ROWID; + +-- models_catalog: cache of /v1/models + /base-models, refreshed hourly. +-- Replaces the model_prices table for display purposes. +-- input/output_price_per_million_micros are from /base-models — display +-- only, not used for billing math. +CREATE TABLE models_catalog ( + id TEXT PRIMARY KEY, + label TEXT NOT NULL, + description TEXT NOT NULL DEFAULT '', + context_window INTEGER, + max_output_tokens INTEGER, + is_chat INTEGER NOT NULL DEFAULT 1, + tier TEXT, + input_price_per_million_micros INTEGER, + output_price_per_million_micros INTEGER, + raw_json TEXT NOT NULL DEFAULT '{}', + refreshed_at INTEGER NOT NULL +) STRICT; + +-- +goose StatementEnd + +-- +goose Down +-- +goose StatementBegin +DROP TABLE IF EXISTS models_catalog; +DROP TABLE IF EXISTS user_daily_allowances; +DROP TABLE IF EXISTS user_daily_spend; +DROP TABLE IF EXISTS pool_key_billing_rows; +DROP TABLE IF EXISTS potluck_requests; + +-- SQLite doesn't support DROP COLUMN before 3.35; the columns added via +-- ALTER TABLE are left in place on rollback. Rolling back to 00005 and +-- forward again is the safe path if you need a clean slate. + +-- +goose StatementEnd diff --git a/server/db/queries/billing_rows.sql b/server/db/queries/billing_rows.sql new file mode 100644 index 0000000..3eca423 --- /dev/null +++ b/server/db/queries/billing_rows.sql @@ -0,0 +1,39 @@ +-- pool_key_billing_rows: ingested pioneer billing log entries. +-- +-- attribution integer enum: +-- 0 = matched (joined to a potluck_requests row) +-- 1 = judge_paired (/llmaj/judge paired to preceding opus call) +-- 2 = owner_fallback (no match, charged to key owner) +-- 3 = duplicate (double-logged by pioneer, not charged) + +-- name: UpsertBillingRow :exec +INSERT INTO pool_key_billing_rows ( + id, pool_key_id, pioneer_created_at, credit_micros, cost_micros, + token_usage, model, endpoint, + attributed_user_id, attribution, is_duplicate, matched_request_id, ingested_at +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +ON CONFLICT(id) DO NOTHING; + +-- name: LatestBillingRowTime :one +-- Most recent pioneer_created_at for a key. +-- COALESCE(SUM(pioneer_created_at*0) + MAX(...)) is a workaround for +-- sqlc's sqlite parser rejecting bare COALESCE(MAX(...), 0). +SELECT COALESCE(SUM(pioneer_created_at * 0) + MAX(pioneer_created_at), 0) +FROM pool_key_billing_rows +WHERE pool_key_id = ?; + +-- name: ListBillingRowsForKeyAfter :many +-- All billing rows for a key after a given timestamp. +SELECT * FROM pool_key_billing_rows +WHERE pool_key_id = ? + AND pioneer_created_at > ? +ORDER BY pioneer_created_at ASC; + +-- name: ListBillingRowsForUserToday :many +-- Billing rows attributed to a user in a time window. +-- Caller filters duplicates and sums in Go. +SELECT * FROM pool_key_billing_rows +WHERE attributed_user_id = ? + AND pioneer_created_at >= ? + AND pioneer_created_at < ? +ORDER BY pioneer_created_at ASC; diff --git a/server/db/queries/daily_spend.sql b/server/db/queries/daily_spend.sql new file mode 100644 index 0000000..c2d4342 --- /dev/null +++ b/server/db/queries/daily_spend.sql @@ -0,0 +1,39 @@ +-- user_daily_spend and user_daily_allowances. + +-- name: UpsertUserDailySpend :exec +INSERT INTO user_daily_spend (user_id, day, shared_spent_micros, private_spent_micros) +VALUES (?, ?, ?, ?) +ON CONFLICT(user_id, day) DO UPDATE SET + shared_spent_micros = shared_spent_micros + excluded.shared_spent_micros, + private_spent_micros = private_spent_micros + excluded.private_spent_micros; + +-- name: GetUserDailySpend :one +SELECT * FROM user_daily_spend WHERE user_id = ? AND day = ?; + +-- name: ListUserDailySpendForDay :many +SELECT * FROM user_daily_spend WHERE day = ?; + +-- name: UpsertUserDailyAllowance :exec +INSERT INTO user_daily_allowances ( + user_id, day, shared_allowance_micros, set_at, set_by_user_id +) VALUES (?, ?, ?, ?, ?) +ON CONFLICT(user_id, day) DO UPDATE SET + shared_allowance_micros = MAX( + excluded.shared_allowance_micros, + user_daily_allowances.shared_allowance_micros + ), + set_at = excluded.set_at, + set_by_user_id = excluded.set_by_user_id; + +-- name: GetUserDailyAllowance :one +SELECT * FROM user_daily_allowances WHERE user_id = ? AND day = ?; + +-- name: ListUserDailyAllowancesForDay :many +SELECT * FROM user_daily_allowances WHERE day = ?; + +-- name: GetLatestRecompute :one +-- Most recent recompute for today, for display in the dashboard. +SELECT set_at, set_by_user_id FROM user_daily_allowances +WHERE day = ? +ORDER BY set_at DESC +LIMIT 1; diff --git a/server/db/queries/models_catalog.sql b/server/db/queries/models_catalog.sql new file mode 100644 index 0000000..0d782d6 --- /dev/null +++ b/server/db/queries/models_catalog.sql @@ -0,0 +1,29 @@ +-- models_catalog: hourly-refreshed cache of /v1/models + /base-models. + +-- name: UpsertModelCatalog :exec +INSERT INTO models_catalog ( + id, label, description, context_window, max_output_tokens, + is_chat, tier, + input_price_per_million_micros, output_price_per_million_micros, + raw_json, refreshed_at +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +ON CONFLICT(id) DO UPDATE SET + label = excluded.label, + description = excluded.description, + context_window = excluded.context_window, + max_output_tokens = excluded.max_output_tokens, + is_chat = excluded.is_chat, + tier = excluded.tier, + input_price_per_million_micros = excluded.input_price_per_million_micros, + output_price_per_million_micros = excluded.output_price_per_million_micros, + raw_json = excluded.raw_json, + refreshed_at = excluded.refreshed_at; + +-- name: ListModelCatalog :many +SELECT * FROM models_catalog +WHERE is_chat = 1 +ORDER BY tier ASC, label ASC; + +-- name: GetModelCatalogRefreshedAt :one +-- Oldest refreshed_at across all models — tells us when the catalog is stale. +SELECT COALESCE(MIN(refreshed_at), 0) FROM models_catalog; diff --git a/server/db/queries/pool_keys_v2.sql b/server/db/queries/pool_keys_v2.sql new file mode 100644 index 0000000..7f34b12 --- /dev/null +++ b/server/db/queries/pool_keys_v2.sql @@ -0,0 +1,76 @@ +-- Pool key v2 queries: health tracking, billing sync, two-budget updates. +-- +-- pioneer_health integer enum: +-- 0 = unknown +-- 1 = healthy +-- 2 = unauthorized + +-- name: UpdatePoolKeyHealth :exec +-- Sets health and snapshots billing info from /plan-info. +-- Pass NULL for optional fields to leave them unchanged. +UPDATE pool_keys SET + pioneer_health = ?, + pioneer_unhealthy_since = ?, + pioneer_team_id = COALESCE(?, pioneer_team_id), + pioneer_payment_plan = COALESCE(?, pioneer_payment_plan), + pioneer_credit_limit_micros = COALESCE(?, pioneer_credit_limit_micros), + pioneer_remaining_micros = COALESCE(?, pioneer_remaining_micros), + today_micros = COALESCE(?, today_micros), + last_billing_sync_at = ? +WHERE id = ?; + +-- name: UpdatePoolKeyLimits :exec +-- Updates the two-budget limits. Server enforces 0 <= shared <= max. +UPDATE pool_keys SET + max_micros = ?, + shared_micros = ? +WHERE id = ? AND user_id = ?; + +-- name: MarkPoolKeyRevoked :exec +UPDATE pool_keys SET + active = 0, + revoked_at = ? +WHERE id = ?; + +-- name: ActivatePoolKey :exec +-- Called by the reconciler when a pending_validation or unauthorized key +-- comes back healthy (pioneer_health=1). +UPDATE pool_keys SET + active = 1, + pending_validation = 0, + pioneer_health = 1, + pioneer_unhealthy_since = NULL +WHERE id = ?; + +-- name: ListKeysNeedingHealthCheck :many +-- Keys the reconciler should probe on each tick. +-- Excludes permanently revoked keys. +SELECT * FROM pool_keys +WHERE revoked_at IS NULL +ORDER BY last_billing_sync_at ASC NULLS FIRST; + +-- name: ListUnhealthyKeysOlderThan :many +-- Keys that have been unauthorized (pioneer_health=2) since before cutoff. +-- Used by the reconciler to trigger permanent revocation after 14 days. +SELECT * FROM pool_keys +WHERE pioneer_health = 2 + AND pioneer_unhealthy_since IS NOT NULL + AND pioneer_unhealthy_since < ? + AND revoked_at IS NULL; + +-- name: PickPoolKeyV2 :one +-- Best active healthy key for a request: +-- active=1, not revoked, not pending validation +-- pioneer_health=1 (healthy) +-- pioneer_remaining_micros > 10,000,000 ($10 buffer = 1000 credits) +-- today_micros < max_micros +-- Lowest today_micros wins; random tiebreak. +SELECT * FROM pool_keys +WHERE active = 1 + AND revoked_at IS NULL + AND pending_validation = 0 + AND pioneer_health = 1 + AND (pioneer_remaining_micros IS NULL OR pioneer_remaining_micros > 10000000) + AND today_micros < max_micros +ORDER BY today_micros ASC, RANDOM() +LIMIT 1; diff --git a/server/db/queries/potluck_requests.sql b/server/db/queries/potluck_requests.sql new file mode 100644 index 0000000..01c0ba4 --- /dev/null +++ b/server/db/queries/potluck_requests.sql @@ -0,0 +1,36 @@ +-- potluck_requests: per-request log for attribution and billing join. + +-- name: CreatePotluckRequest :one +INSERT INTO potluck_requests ( + id, user_id, api_key_id, pool_key_id, surface, model, started_at, status +) VALUES (?, ?, ?, ?, ?, ?, ?, 'pending') +RETURNING *; + +-- name: FinishPotluckRequest :exec +UPDATE potluck_requests SET + finished_at = ?, + prompt_tokens = ?, + completion_tokens = ?, + total_tokens = ?, + status = ? +WHERE id = ?; + +-- name: CancelPotluckRequest :exec +UPDATE potluck_requests SET + finished_at = ?, + status = 'canceled' +WHERE id = ? AND status = 'pending'; + +-- name: ListUnmatchedRequestsForKey :many +-- Requests that finished in a given time window with no billing row matched yet. +-- Used by the reconciler attribution pass. +SELECT pr.* FROM potluck_requests pr +WHERE pr.pool_key_id = ? + AND pr.finished_at >= ? + AND pr.finished_at <= ? + AND pr.status = 'done' + AND pr.id NOT IN ( + SELECT pkbr.matched_request_id FROM pool_key_billing_rows pkbr + WHERE pkbr.matched_request_id IS NOT NULL + ) +ORDER BY pr.finished_at ASC; diff --git a/server/internal/api/v1/chat.go b/server/internal/api/v1/chat.go index 512875f..8be002b 100644 --- a/server/internal/api/v1/chat.go +++ b/server/internal/api/v1/chat.go @@ -3,26 +3,26 @@ package v1 import ( "bytes" "context" + "database/sql" "encoding/json" "io" "net/http" + "time" + "github.com/google/uuid" + + apimw "github.com/taciturnaxolotl/potluck/internal/api/middleware" + "github.com/taciturnaxolotl/potluck/internal/pool" "github.com/taciturnaxolotl/potluck/internal/provider" + "github.com/taciturnaxolotl/potluck/internal/store" ) -// handleChatCompletions proxies POST /v1/chat/completions to pioneer -// with minimal fuss. Streaming requests stay streaming; non-streaming -// responses are buffered (small payloads, OpenAI shape). -// -// Cancellation semantics differ from the /api/* surface: here the -// upstream is bound to the request context. Client disconnect ➜ upstream -// canceled ➜ no spend recorded for tokens we won't deliver. This is the -// right choice for stateless API clients that aren't refreshing tabs. +// handleChatCompletions proxies POST /v1/chat/completions to pioneer. +// Streaming requests stay streaming; non-streaming responses are buffered. // -// Spend recording is intentionally NOT in this stub yet — see -// design/public-api.md. Wiring it up requires settling against -// stream_options.include_usage on the streaming path and the response's -// own usage block on the non-streaming path. +// Cancellation semantics: the upstream is bound to the request context. +// Client disconnect → upstream canceled → no spend for tokens we didn't deliver. +// This is correct for stateless API clients (not refreshing tabs). func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) { body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, 1<<20)) if err != nil { @@ -31,9 +31,6 @@ func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) { } defer r.Body.Close() - // Peek to decide stream vs JSON without unmarshalling the whole body - // into a typed struct (which we deliberately avoid — pioneer surfaces - // fields we may not know about and clients expect those passed through). var probe struct { Stream bool `json:"stream"` Model string `json:"model"` @@ -44,21 +41,43 @@ func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) { } if probe.Stream { - s.streamCompletion(w, r, body) + s.streamCompletion(w, r, body, probe.Model) return } - s.bufferedCompletion(w, r, body) + s.bufferedCompletion(w, r, body, probe.Model) } -// bufferedCompletion handles a non-streaming chat completion: forward the -// body, return whatever pioneer returns. We do NOT re-shape the response. -func (s *Server) bufferedCompletion(w http.ResponseWriter, r *http.Request, body []byte) { +// bufferedCompletion handles a non-streaming chat completion. +func (s *Server) bufferedCompletion(w http.ResponseWriter, r *http.Request, body []byte, model string) { sel, err := s.Pool.Pick(r.Context()) if err != nil { writeError(w, http.StatusServiceUnavailable, "no_pool_keys", "no active pool keys available") return } + u, _ := apimw.UserFromContext(r.Context()) + apiKey, _ := apimw.APIKeyFromContext(r.Context()) + + // Write request log row immediately. + reqID := uuid.NewString() + poolKeyID := sql.NullString{String: sel.KeyID(), Valid: sel.KeyID() != ""} + apiKeyID := sql.NullString{} + if apiKey != nil { + apiKeyID = sql.NullString{String: apiKey.ID, Valid: true} + } + startedAt := time.Now().Unix() + if u != nil { + _, _ = s.Q.CreatePotluckRequest(r.Context(), store.CreatePotluckRequestParams{ + ID: reqID, + UserID: u.ID, + ApiKeyID: apiKeyID, + PoolKeyID: poolKeyID, + Surface: "v1", + Model: model, + StartedAt: startedAt, + }) + } + req, err := http.NewRequestWithContext(r.Context(), http.MethodPost, s.Provider.BaseURL+"/v1/chat/completions", bytes.NewReader(body)) if err != nil { @@ -71,35 +90,74 @@ func (s *Server) bufferedCompletion(w http.ResponseWriter, r *http.Request, body resp, err := s.Provider.HTTP.Do(req) if err != nil { writeError(w, http.StatusBadGateway, "provider_down", err.Error()) + if u != nil { + go finishRequest(s.Q, reqID, 0, 0, 0, "error") + } return } defer resp.Body.Close() + // Buffer the response so we can parse usage before returning. + respBody, _ := io.ReadAll(resp.Body) w.Header().Set("Content-Type", "application/json") w.WriteHeader(resp.StatusCode) - _, _ = io.Copy(w, resp.Body) + _, _ = w.Write(respBody) - // Record spend asynchronously — don't hold the client on a DB write. - // Non-streaming: we don't have token counts here yet, record 0 to - // increment request_count and last_used_at. TODO: parse response usage. + // Parse usage from the response body and settle asynchronously. go func() { + var respJSON struct { + Usage *provider.Usage `json:"usage"` + } + _ = json.Unmarshal(respBody, &respJSON) + var prompt, completion, total int64 + if respJSON.Usage != nil { + prompt = int64(respJSON.Usage.PromptTokens) + completion = int64(respJSON.Usage.CompletionTokens) + total = int64(respJSON.Usage.TotalTokens) + } + status := "done" + if resp.StatusCode/100 != 2 { + status = "error" + } + if u != nil { + finishRequest(s.Q, reqID, prompt, completion, total, status) + } _ = s.Pool.RecordSpend(context.Background(), sel, 0) }() } -// streamCompletion forwards an SSE chat completion straight through. We -// use provider.StreamChat for the chunk parser (so we can settle spend at -// the end) but the bytes the client sees are pioneer's verbatim where -// possible. -func (s *Server) streamCompletion(w http.ResponseWriter, r *http.Request, body []byte) { +// streamCompletion forwards an SSE chat completion straight through. +func (s *Server) streamCompletion(w http.ResponseWriter, r *http.Request, body []byte, model string) { sel, err := s.Pool.Pick(r.Context()) if err != nil { writeError(w, http.StatusServiceUnavailable, "no_pool_keys", "no active pool keys available") return } - // Decode just enough to ensure stream_options.include_usage is on, so - // we can settle accurately. Re-marshal and forward. + u, _ := apimw.UserFromContext(r.Context()) + apiKey, _ := apimw.APIKeyFromContext(r.Context()) + + // Write request log row immediately. + reqID := uuid.NewString() + poolKeyID := sql.NullString{String: sel.KeyID(), Valid: sel.KeyID() != ""} + apiKeyID := sql.NullString{} + if apiKey != nil { + apiKeyID = sql.NullString{String: apiKey.ID, Valid: true} + } + startedAt := time.Now().Unix() + if u != nil { + _, _ = s.Q.CreatePotluckRequest(r.Context(), store.CreatePotluckRequestParams{ + ID: reqID, + UserID: u.ID, + ApiKeyID: apiKeyID, + PoolKeyID: poolKeyID, + Surface: "v1", + Model: model, + StartedAt: startedAt, + }) + } + + // Ensure stream_options.include_usage is on for accurate settlement. var req map[string]any if err := json.Unmarshal(body, &req); err != nil { writeError(w, http.StatusBadRequest, "invalid_request", err.Error()) @@ -110,7 +168,6 @@ func (s *Server) streamCompletion(w http.ResponseWriter, r *http.Request, body [ req["stream_options"] = map[string]any{"include_usage": true} } - // Use a per-request client copy with the selected pool key. pc := &provider.Client{ BaseURL: s.Provider.BaseURL, APIKey: sel.APIKey(), @@ -124,6 +181,9 @@ func (s *Server) streamCompletion(w http.ResponseWriter, r *http.Request, body [ }) if err != nil { writeError(w, http.StatusBadGateway, "provider_down", err.Error()) + if u != nil { + go finishRequest(s.Q, reqID, 0, 0, 0, "error") + } return } @@ -133,23 +193,29 @@ func (s *Server) streamCompletion(w http.ResponseWriter, r *http.Request, body [ w.WriteHeader(http.StatusOK) flusher, _ := w.(http.Flusher) - var totalMicros int64 + var usage *provider.Usage for { select { case <-r.Context().Done(): + if u != nil { + go finishRequest(s.Q, reqID, 0, 0, 0, "canceled") + } return case ch, ok := <-chunks: if !ok { _, _ = w.Write([]byte("data: [DONE]\n\n")) - go func() { _ = s.Pool.RecordSpend(context.Background(), sel, totalMicros) }() + go settle(s.Q, s.Pool, sel, reqID, usage, u) return } + if ch.Usage != nil { + usage = ch.Usage + } if ch.Done { _, _ = w.Write([]byte("data: [DONE]\n\n")) if flusher != nil { flusher.Flush() } - go func() { _ = s.Pool.RecordSpend(context.Background(), sel, totalMicros) }() + go settle(s.Q, s.Pool, sel, reqID, usage, u) return } _, _ = w.Write([]byte("data: ")) @@ -161,12 +227,42 @@ func (s *Server) streamCompletion(w http.ResponseWriter, r *http.Request, body [ case e := <-errs: if e != nil { writeError(w, http.StatusBadGateway, "provider_error", e.Error()) + if u != nil { + go finishRequest(s.Q, reqID, 0, 0, 0, "error") + } return } } } } +// settle fires off the post-stream DB writes. Runs in a goroutine. +func settle(q *store.Queries, poolMgr *pool.Manager, sel *pool.Selection, reqID string, usage *provider.Usage, u *store.User) { + var prompt, completion, total int64 + if usage != nil { + prompt = int64(usage.PromptTokens) + completion = int64(usage.CompletionTokens) + total = int64(usage.TotalTokens) + } + if u != nil { + finishRequest(q, reqID, prompt, completion, total, "done") + } + _ = poolMgr.RecordSpend(context.Background(), sel, 0) +} + +// finishRequest updates the potluck_requests row after the upstream call ends. +func finishRequest(q *store.Queries, reqID string, prompt, completion, total int64, status string) { + now := time.Now().Unix() + _ = q.FinishPotluckRequest(context.Background(), store.FinishPotluckRequestParams{ + FinishedAt: sql.NullInt64{Int64: now, Valid: true}, + PromptTokens: sql.NullInt64{Int64: prompt, Valid: prompt > 0}, + CompletionTokens: sql.NullInt64{Int64: completion, Valid: completion > 0}, + TotalTokens: sql.NullInt64{Int64: total, Valid: total > 0}, + Status: status, + ID: reqID, + }) +} + // asString safely extracts a string from a map[string]any without panicking. func asString(v any) string { s, _ := v.(string) @@ -174,8 +270,7 @@ func asString(v any) string { } // messagesFromMap turns the JSON-decoded `messages` array into the typed -// slice provider.StreamChat wants. We accept map-shaped messages; anything -// else is left to pioneer to reject. +// slice provider.StreamChat wants. func messagesFromMap(v any) []provider.ChatMessage { arr, ok := v.([]any) if !ok { diff --git a/server/internal/api/web/models.go b/server/internal/api/web/models.go index e9e16f5..f368752 100644 --- a/server/internal/api/web/models.go +++ b/server/internal/api/web/models.go @@ -166,7 +166,14 @@ func fetchMergedModels(ctx context.Context, apiKey string) ([]mergedModel, error } func (s *Server) handleListModels(w http.ResponseWriter, r *http.Request) { - models, err := fetchMergedModels(r.Context(), s.PioneerAPIKey) + // Use a pool key for the model fetch. If the pool is empty we still + // try with an empty key — pioneer's /v1/models requires auth but the + // error is graceful. + apiKey := "" + if sel, err := s.Pool.Pick(r.Context()); err == nil { + apiKey = sel.APIKey() + } + models, err := fetchMergedModels(r.Context(), apiKey) if err != nil { writeErr(w, 502, "upstream_error", "could not fetch model list from provider") return diff --git a/server/internal/api/web/null.go b/server/internal/api/web/null.go index 49d50fb..d5e71e9 100644 --- a/server/internal/api/web/null.go +++ b/server/internal/api/web/null.go @@ -11,6 +11,20 @@ func sqlNullInt64Ptr(v *int64) sql.NullInt64 { return sql.NullInt64{Int64: *v, Valid: true} } +func nullStrWeb(s string) sql.NullString { + if s == "" { + return sql.NullString{} + } + return sql.NullString{String: s, Valid: true} +} + +func nullIntWeb(v int64) sql.NullInt64 { + if v == 0 { + return sql.NullInt64{} + } + return sql.NullInt64{Int64: v, Valid: true} +} + // toInt64 unboxes the interface{} sqlc emits for SUM/COALESCE aggregates. func toInt64(v any) int64 { switch x := v.(type) { diff --git a/server/internal/api/web/pool_keys.go b/server/internal/api/web/pool_keys.go index 73c1172..e491815 100644 --- a/server/internal/api/web/pool_keys.go +++ b/server/internal/api/web/pool_keys.go @@ -1,9 +1,12 @@ package web import ( + "bytes" "context" + "database/sql" "encoding/json" "fmt" + "io" "net/http" "strings" "time" @@ -27,19 +30,28 @@ func (s *Server) handleListPoolKeys(w http.ResponseWriter, r *http.Request) { out := make([]map[string]any, 0, len(rows)) for _, k := range rows { out = append(out, map[string]any{ - "id": k.ID, - "user_id": k.UserID, - "label": k.Label, - "active": k.Active == 1, - "daily_limit_micros": k.DailyLimitMicros, - "today_micros": k.TodayMicros, - "total_micros": k.TotalMicros, - "request_count": k.RequestCount, - "created_at": k.CreatedAt, - "last_used_at": k.LastUsedAt.Int64, - "owner_name": k.OwnerName, - "owner_email": k.OwnerEmail, - "mine": k.UserID == u.ID, + "id": k.ID, + "user_id": k.UserID, + "label": k.Label, + "active": k.Active == 1, + "max_micros": k.MaxMicros, + "shared_micros": k.SharedMicros, + "private_micros": k.MaxMicros - k.SharedMicros, + "today_micros": k.TodayMicros, + "total_micros": k.TotalMicros, + "request_count": k.RequestCount, + "pioneer_health": k.PioneerHealth, + "pioneer_team_id": k.PioneerTeamID.String, + "pioneer_payment_plan": k.PioneerPaymentPlan.String, + "pioneer_credit_limit_micros": k.PioneerCreditLimitMicros.Int64, + "pioneer_remaining_micros": k.PioneerRemainingMicros.Int64, + "pending_validation": k.PendingValidation == 1, + "created_at": k.CreatedAt, + "last_used_at": k.LastUsedAt.Int64, + "last_billing_sync_at": k.LastBillingSyncAt.Int64, + "owner_name": k.OwnerName, + "owner_email": k.OwnerEmail, + "mine": k.UserID == u.ID, }) } writeJSON(w, 200, out) @@ -55,61 +67,103 @@ const pioneerBillingTimeseriesURL = "https://api.pioneer.ai/billing/usage/timese // pioneerBillingResult holds the result of probing pioneer's billing timeseries. type pioneerBillingResult struct { - TodayMicros int64 // today's spend in micros (credits × 1 ≈ micros) + TodayMicros int64 // today's spend in micros + RemainingMicros int64 // remaining credits in micros + CreditLimitMicros int64 // credit limit in micros + TeamID string + PaymentPlan string + HTTP401 bool // key exhausted or invalid — don't reject, save as pending + HTTP503 bool // pioneer auth down — transient } -// probePioneerBilling calls pioneer's billing timeseries to validate the key -// and fetch today's spend. Returns an error if the key is rejected. -// pioneer credits are treated as micros (1 credit ≈ 1 USD micro). +// probePioneerBilling calls pioneer's billing endpoints to validate the key +// and fetch today's spend. Never returns error for 401/503 — use the flags. func probePioneerBilling(ctx context.Context, apiKey string) (pioneerBillingResult, error) { - ctx, cancel := context.WithTimeout(ctx, 8*time.Second) + ctx, cancel := context.WithTimeout(ctx, 12*time.Second) defer cancel() - todayUTC := time.Now().UTC().Format("2006-01-02") - - req, err := http.NewRequestWithContext(ctx, http.MethodGet, - pioneerBillingTimeseriesURL+"?full_history=true&interval_minutes=1440", nil) - if err != nil { - return pioneerBillingResult{}, fmt.Errorf("could not build billing request: %w", err) + do := func(url string) ([]byte, int, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, 0, err + } + req.Header.Set("Authorization", "Bearer "+apiKey) + resp, err := http.DefaultClient.Do(req) + if err != nil { + return nil, 0, err + } + defer resp.Body.Close() + b, _ := io.ReadAll(resp.Body) + return b, resp.StatusCode, nil } - req.Header.Set("Authorization", "Bearer "+apiKey) - resp, err := http.DefaultClient.Do(req) + body, status, err := do(pioneerBillingTimeseriesURL + "?full_history=false&interval_minutes=1440") if err != nil { return pioneerBillingResult{}, fmt.Errorf("billing request failed: %w", err) } - defer resp.Body.Close() - - switch resp.StatusCode { - case http.StatusUnauthorized, http.StatusForbidden: - return pioneerBillingResult{}, fmt.Errorf("pioneer rejected the key (HTTP %d) — double-check it", resp.StatusCode) + switch status { + case http.StatusUnauthorized: + return pioneerBillingResult{HTTP401: true}, nil + case http.StatusServiceUnavailable: + return pioneerBillingResult{HTTP503: true}, nil } - if resp.StatusCode/100 != 2 { - return pioneerBillingResult{}, fmt.Errorf("pioneer returned HTTP %d during validation", resp.StatusCode) + if status/100 != 2 { + return pioneerBillingResult{}, fmt.Errorf("pioneer returned HTTP %d during validation", status) } - var body struct { + var tsBody struct { Points []struct { BucketDate string `json:"bucket_date"` TotalCredits float64 `json:"total_credits"` } `json:"points"` } - if err := json.NewDecoder(resp.Body).Decode(&body); err != nil { + if err := json.NewDecoder(bytes.NewReader(body)).Decode(&tsBody); err != nil { return pioneerBillingResult{}, fmt.Errorf("could not decode billing response: %w", err) } + // Fetch plan-info for credit limit, remaining, team_id, payment_plan. + body2, status2, err2 := do("https://api.pioneer.ai/billing/plan-info") + body3, status3, err3 := do("https://api.pioneer.ai/billing/billing-status") + + var planInfo struct { + PaymentPlan string `json:"payment_plan"` + CreditLimit float64 `json:"credit_limit"` + RemainingCredits float64 `json:"remaining_credits"` + } + if err2 == nil && status2 == http.StatusOK { + _ = json.NewDecoder(bytes.NewReader(body2)).Decode(&planInfo) + } + + var billingStatus struct { + TeamID string `json:"team_id"` + } + if err3 == nil && status3 == http.StatusOK { + _ = json.NewDecoder(bytes.NewReader(body3)).Decode(&billingStatus) + } + + todayUTC := time.Now().UTC().Format("2006-01-02") var todayMicros int64 - for _, p := range body.Points { + for _, p := range tsBody.Points { if p.BucketDate == todayUTC { - // 1 pioneer credit = $0.01 = 10,000 micros todayMicros = int64(p.TotalCredits * 10_000) break } } - return pioneerBillingResult{TodayMicros: todayMicros}, nil + + return pioneerBillingResult{ + TodayMicros: todayMicros, + RemainingMicros: int64(planInfo.RemainingCredits * 10_000), + CreditLimitMicros: int64(planInfo.CreditLimit * 10_000), + TeamID: billingStatus.TeamID, + PaymentPlan: planInfo.PaymentPlan, + }, nil } // handleAddPoolKey validates a key against pioneer, then encrypts and stores it. +// If pioneer returns 401 (key exhausted or not yet active), we save it as +// pending_validation and let the reconciler activate it when it comes back. +// If pioneer returns 503 (auth service down), we also save as pending. +// Only pro-plan keys are accepted. func (s *Server) handleAddPoolKey(w http.ResponseWriter, r *http.Request) { u, _ := currentUser(r) var req addPoolKeyReq @@ -124,6 +178,22 @@ func (s *Server) handleAddPoolKey(w http.ResponseWriter, r *http.Request) { return } + // Determine whether we can activate immediately. + pendingValidation := int64(0) + var pendingReason string + switch { + case billing.HTTP401: + pendingValidation = 1 + pendingReason = "pioneer returned 401 — key may be exhausted or not yet active; we'll retry automatically" + case billing.HTTP503: + pendingValidation = 1 + pendingReason = "pioneer auth service is temporarily down; we'll retry automatically" + case billing.PaymentPlan != "" && billing.PaymentPlan != pool.RequiredPaymentPlan: + writeErr(w, 422, "invalid_plan", + fmt.Sprintf("only pro-plan pioneer keys are supported (got %q)", billing.PaymentPlan)) + return + } + fingerprint := pool.Fingerprint(req.APIKey) ciphertext, err := s.Pool.Encrypt(req.APIKey) if err != nil { @@ -131,28 +201,23 @@ func (s *Server) handleAddPoolKey(w http.ResponseWriter, r *http.Request) { return } - dailyLimit := int64(1_000_000_000) // $1000 default (1 USD = 1_000_000 micros, $1000 = 1e9) + maxMicros := int64(1_000_000_000) // $1000 default + sharedMicros := int64(1_000_000_000) if req.DailyLimitMicros != nil && *req.DailyLimitMicros > 0 { - dailyLimit = *req.DailyLimitMicros + maxMicros = *req.DailyLimitMicros + sharedMicros = *req.DailyLimitMicros } + now := time.Now().Unix() row, err := s.Q.CreatePoolKey(r.Context(), store.CreatePoolKeyParams{ ID: uuid.NewString(), UserID: u.ID, Label: req.Label, KeyCiphertext: ciphertext, KeyFingerprint: fingerprint, - DailyLimitMicros: dailyLimit, - CreatedAt: time.Now().Unix(), + DailyLimitMicros: maxMicros, + CreatedAt: now, }) - if err == nil && billing.TodayMicros > 0 { - // Seed today's real spend from pioneer immediately. - _ = s.Q.SyncTodaySpend(r.Context(), store.SyncTodaySpendParams{ - TodayDate: int64(time.Now().UTC().Unix() / 86400), - TodayMicros: billing.TodayMicros, - ID: row.ID, - }) - } if err != nil { if isUniqueConstraintErr(err) { writeErr(w, 409, "duplicate_key", "this API key is already in the pool") @@ -162,18 +227,61 @@ func (s *Server) handleAddPoolKey(w http.ResponseWriter, r *http.Request) { return } - writeJSON(w, 201, map[string]any{ - "id": row.ID, - "user_id": row.UserID, - "label": row.Label, - "active": row.Active == 1, - "daily_limit_micros": row.DailyLimitMicros, - "today_micros": row.TodayMicros, - "total_micros": row.TotalMicros, - "request_count": row.RequestCount, - "created_at": row.CreatedAt, - "mine": true, + // Backfill the v2 columns not in CreatePoolKey (additive migration). + _ = s.Q.UpdatePoolKeyLimits(r.Context(), store.UpdatePoolKeyLimitsParams{ + MaxMicros: maxMicros, + SharedMicros: sharedMicros, + ID: row.ID, + UserID: u.ID, }) + + // Set health and billing snapshot. + health := pool.HealthUnknown + if pendingValidation == 0 { + health = pool.HealthHealthy + } + syncAt := sql.NullInt64{} + if pendingValidation == 0 { + syncAt = sql.NullInt64{Int64: now, Valid: true} + } + _ = s.Q.UpdatePoolKeyHealth(r.Context(), store.UpdatePoolKeyHealthParams{ + PioneerHealth: health, + PioneerUnhealthySince: sql.NullInt64{}, + PioneerTeamID: nullStrWeb(billing.TeamID), + PioneerPaymentPlan: nullStrWeb(billing.PaymentPlan), + PioneerCreditLimitMicros: nullIntWeb(billing.CreditLimitMicros), + PioneerRemainingMicros: nullIntWeb(billing.RemainingMicros), + TodayMicros: billing.TodayMicros, + LastBillingSyncAt: syncAt, + ID: row.ID, + }) + + // Seed today_micros if we got real data. + if billing.TodayMicros > 0 && pendingValidation == 0 { + _ = s.Q.SyncTodaySpend(r.Context(), store.SyncTodaySpendParams{ + TodayDate: int64(time.Now().UTC().Unix() / 86400), + TodayMicros: billing.TodayMicros, + ID: row.ID, + }) + } + + resp := map[string]any{ + "id": row.ID, + "user_id": row.UserID, + "label": row.Label, + "active": row.Active == 1 && pendingValidation == 0, + "max_micros": maxMicros, + "shared_micros": sharedMicros, + "today_micros": billing.TodayMicros, + "pioneer_health": health, + "pending_validation": pendingValidation == 1, + "created_at": row.CreatedAt, + "mine": true, + } + if pendingReason != "" { + resp["pending_reason"] = pendingReason + } + writeJSON(w, 201, resp) } // handleSetPoolKeyActive toggles the active state of a pool key. @@ -240,13 +348,14 @@ func (s *Server) handleUpdatePoolKeyLabel(w http.ResponseWriter, r *http.Request w.WriteHeader(204) } -// handleUpdatePoolKeyLimit adjusts the daily share limit for a key. -// Value is clamped server-side to [$100, $1000] in micros. -func (s *Server) handleUpdatePoolKeyLimit(w http.ResponseWriter, r *http.Request) { +// handleUpdatePoolKeyLimits sets both max_micros and shared_micros. +// Server enforces 0 <= shared <= max, both clamped to [$100, $1000]. +func (s *Server) handleUpdatePoolKeyLimits(w http.ResponseWriter, r *http.Request) { u, _ := currentUser(r) id := chi.URLParam(r, "id") var body struct { - DailyLimitMicros int64 `json:"daily_limit_micros"` + MaxMicros int64 `json:"max_micros"` + SharedMicros int64 `json:"shared_micros"` } if err := json.NewDecoder(r.Body).Decode(&body); err != nil { writeErr(w, 400, "invalid_request", err.Error()) @@ -254,16 +363,23 @@ func (s *Server) handleUpdatePoolKeyLimit(w http.ResponseWriter, r *http.Request } const minMicros = 100_000_000 // $100 const maxMicros = 1_000_000_000 // $1000 - if body.DailyLimitMicros < minMicros { - body.DailyLimitMicros = minMicros + if body.MaxMicros < minMicros { + body.MaxMicros = minMicros } - if body.DailyLimitMicros > maxMicros { - body.DailyLimitMicros = maxMicros + if body.MaxMicros > maxMicros { + body.MaxMicros = maxMicros } - if err := s.Q.UpdatePoolKeyLimit(r.Context(), store.UpdatePoolKeyLimitParams{ - DailyLimitMicros: body.DailyLimitMicros, - ID: id, - UserID: u.ID, + if body.SharedMicros < 0 { + body.SharedMicros = 0 + } + if body.SharedMicros > body.MaxMicros { + body.SharedMicros = body.MaxMicros + } + if err := s.Q.UpdatePoolKeyLimits(r.Context(), store.UpdatePoolKeyLimitsParams{ + MaxMicros: body.MaxMicros, + SharedMicros: body.SharedMicros, + ID: id, + UserID: u.ID, }); err != nil { writeErr(w, 500, "internal", err.Error()) return diff --git a/server/internal/api/web/web.go b/server/internal/api/web/web.go index 3bca695..c4371f8 100644 --- a/server/internal/api/web/web.go +++ b/server/internal/api/web/web.go @@ -27,12 +27,11 @@ import ( // Server bundles the deps the web handlers need. type Server struct { - Q *store.Queries - Auth *auth.Service - Ledger *ledger.Service - Hub *stream.Hub - Pool *pool.Manager - PioneerAPIKey string + Q *store.Queries + Auth *auth.Service + Ledger *ledger.Service + Hub *stream.Hub + Pool *pool.Manager } // Mount registers /api/* routes on r. The caller wraps with cookie-auth @@ -61,7 +60,7 @@ func (s *Server) Mount(r chi.Router) { r.Post("/pool-keys", s.handleAddPoolKey) r.Patch("/pool-keys/{id}/active", s.handleSetPoolKeyActive) r.Patch("/pool-keys/{id}/label", s.handleUpdatePoolKeyLabel) - r.Patch("/pool-keys/{id}/limit", s.handleUpdatePoolKeyLimit) + r.Patch("/pool-keys/{id}/limits", s.handleUpdatePoolKeyLimits) r.Post("/pool-keys/{id}/sync", s.handleSyncPoolKeySpend) r.Delete("/pool-keys/{id}", s.handleDeletePoolKey) diff --git a/server/internal/migrations/files/00006_pool_keys_v2.sql b/server/internal/migrations/files/00006_pool_keys_v2.sql new file mode 100644 index 0000000..34f9d35 --- /dev/null +++ b/server/internal/migrations/files/00006_pool_keys_v2.sql @@ -0,0 +1,140 @@ +-- +goose Up +-- +goose StatementBegin + +-- Extend pool_keys with two-budget model, pioneer health tracking, and +-- pending-validation support. All changes are additive; existing rows +-- and queries keep working unchanged until the follow-up migration drops +-- daily_limit_micros. +-- +-- Two-budget model: +-- max_micros = owner's daily ceiling (replaces daily_limit_micros) +-- shared_micros = portion donated to the shared pool (≤ max_micros) +-- private = max_micros - shared_micros (reserved for owner only) +-- +-- Pioneer health: +-- pioneer_health = 'unknown' | 'healthy' | 'unauthorized' +-- A 401 from pioneer means exhausted OR revoked — we can't tell which. +-- Mark unauthorized, retry daily. After 14 days consecutive, set revoked_at. +-- A 503 from pioneer means their auth service is down — don't touch health. +-- +-- Only pro-plan keys are accepted (credit_limit = 100,000 credits = $1000/day). + +ALTER TABLE pool_keys ADD COLUMN max_micros INTEGER NOT NULL DEFAULT 1000000000; +ALTER TABLE pool_keys ADD COLUMN shared_micros INTEGER NOT NULL DEFAULT 1000000000; +ALTER TABLE pool_keys ADD COLUMN pioneer_team_id TEXT; +ALTER TABLE pool_keys ADD COLUMN pioneer_payment_plan TEXT; +ALTER TABLE pool_keys ADD COLUMN pioneer_credit_limit_micros INTEGER; +ALTER TABLE pool_keys ADD COLUMN pioneer_remaining_micros INTEGER; +ALTER TABLE pool_keys ADD COLUMN pioneer_health INTEGER NOT NULL DEFAULT 0; + -- 0=unknown 1=healthy 2=unauthorized +ALTER TABLE pool_keys ADD COLUMN pioneer_unhealthy_since INTEGER; +ALTER TABLE pool_keys ADD COLUMN pending_validation INTEGER NOT NULL DEFAULT 0; +ALTER TABLE pool_keys ADD COLUMN last_billing_sync_at INTEGER; +ALTER TABLE pool_keys ADD COLUMN revoked_at INTEGER; + +-- Backfill: seed max/shared from the old daily_limit_micros. +UPDATE pool_keys SET + max_micros = daily_limit_micros, + shared_micros = daily_limit_micros; + +-- potluck_requests: one row per chat completion we proxy, written at +-- request start and updated at finish. The reconciler matches pioneer +-- billing rows against this table to attribute cost to the right user. +CREATE TABLE potluck_requests ( + id TEXT PRIMARY KEY, -- our uuid + user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE, + api_key_id TEXT REFERENCES api_keys(id) ON DELETE SET NULL, + pool_key_id TEXT REFERENCES pool_keys(id) ON DELETE SET NULL, + surface TEXT NOT NULL, -- 'web' | 'v1' + model TEXT NOT NULL, + started_at INTEGER NOT NULL, -- unix seconds + finished_at INTEGER, -- NULL while in flight + prompt_tokens INTEGER, + completion_tokens INTEGER, + total_tokens INTEGER, + status TEXT NOT NULL DEFAULT 'pending', -- 'pending' | 'done' | 'error' | 'canceled' + error_code TEXT +) STRICT; +CREATE INDEX potluck_requests_by_user ON potluck_requests(user_id, started_at DESC); +CREATE INDEX potluck_requests_by_key ON potluck_requests(pool_key_id, finished_at); + +-- pool_key_billing_rows: one row per pioneer billing log entry ingested. +-- Idempotent on pioneer's id. Attribution tracks which potluck user +-- caused the spend (NULL = key owner, for off-platform or judge fallback). +CREATE TABLE pool_key_billing_rows ( + id TEXT PRIMARY KEY, -- pioneer's billing row UUID + pool_key_id TEXT NOT NULL REFERENCES pool_keys(id) ON DELETE CASCADE, + pioneer_created_at INTEGER NOT NULL, -- unix seconds + credit_micros INTEGER NOT NULL, -- credit_usage * 10000 + cost_micros INTEGER NOT NULL, -- cost_usd * 1000000 + token_usage INTEGER NOT NULL, + model TEXT NOT NULL, + endpoint TEXT NOT NULL, + attributed_user_id TEXT REFERENCES users(id), -- NULL = key owner + -- attribution: 0=matched 1=judge_paired 2=owner_fallback 3=duplicate + attribution INTEGER NOT NULL, + is_duplicate INTEGER NOT NULL DEFAULT 0, -- 1 when attribution=3 + matched_request_id TEXT REFERENCES potluck_requests(id), + ingested_at INTEGER NOT NULL +) STRICT; +CREATE INDEX pool_key_billing_by_key ON pool_key_billing_rows(pool_key_id, pioneer_created_at DESC); +CREATE INDEX pool_key_billing_by_user ON pool_key_billing_rows(attributed_user_id, pioneer_created_at DESC); + +-- user_daily_spend: per-user-per-UTC-day spend split into shared vs +-- private. Updated by the reconciler in the same transaction as +-- pool_key_billing_rows inserts. +-- day = unix_seconds / 86400. +CREATE TABLE user_daily_spend ( + user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE, + day INTEGER NOT NULL, + shared_spent_micros INTEGER NOT NULL DEFAULT 0, + private_spent_micros INTEGER NOT NULL DEFAULT 0, + PRIMARY KEY (user_id, day) +) STRICT, WITHOUT ROWID; + +-- user_daily_allowances: how much of the shared pool each user may +-- spend today. Set by the recompute button. Allowances only ever grow +-- on a recompute (no claw-back). If no row exists for today, the gate +-- falls back to a fair-share calculation on the fly. +CREATE TABLE user_daily_allowances ( + user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE, + day INTEGER NOT NULL, + shared_allowance_micros INTEGER NOT NULL, + set_at INTEGER NOT NULL, + set_by_user_id TEXT NOT NULL REFERENCES users(id), + PRIMARY KEY (user_id, day) +) STRICT, WITHOUT ROWID; + +-- models_catalog: cache of /v1/models + /base-models, refreshed hourly. +-- Replaces the model_prices table for display purposes. +-- input/output_price_per_million_micros are from /base-models — display +-- only, not used for billing math. +CREATE TABLE models_catalog ( + id TEXT PRIMARY KEY, + label TEXT NOT NULL, + description TEXT NOT NULL DEFAULT '', + context_window INTEGER, + max_output_tokens INTEGER, + is_chat INTEGER NOT NULL DEFAULT 1, + tier TEXT, + input_price_per_million_micros INTEGER, + output_price_per_million_micros INTEGER, + raw_json TEXT NOT NULL DEFAULT '{}', + refreshed_at INTEGER NOT NULL +) STRICT; + +-- +goose StatementEnd + +-- +goose Down +-- +goose StatementBegin +DROP TABLE IF EXISTS models_catalog; +DROP TABLE IF EXISTS user_daily_allowances; +DROP TABLE IF EXISTS user_daily_spend; +DROP TABLE IF EXISTS pool_key_billing_rows; +DROP TABLE IF EXISTS potluck_requests; + +-- SQLite doesn't support DROP COLUMN before 3.35; the columns added via +-- ALTER TABLE are left in place on rollback. Rolling back to 00005 and +-- forward again is the safe path if you need a clean slate. + +-- +goose StatementEnd diff --git a/server/internal/pool/pool.go b/server/internal/pool/pool.go index f9a6842..3841958 100644 --- a/server/internal/pool/pool.go +++ b/server/internal/pool/pool.go @@ -32,9 +32,8 @@ var ErrNoKeys = errors.New("pool: no active keys available") // Manager is the pool key manager. Embed in handler structs or pass as a dep. type Manager struct { - q *store.Queries - secret []byte // 32 bytes for AES-256; nil = encryption disabled - fallbackKey string // used when pool is empty (PIONEER_API_KEY) + q *store.Queries + secret []byte // 32 bytes for AES-256; nil = encryption disabled } // New constructs a Manager. @@ -44,8 +43,8 @@ type Manager struct { // encryption — only safe for dev without real keys in the DB. // - fallbackKey: the static PIONEER_API_KEY; used when the pool has no // active keys. Pass "" to disable fallback. -func New(q *store.Queries, secretHex string, fallbackKey string) (*Manager, error) { - m := &Manager{q: q, fallbackKey: fallbackKey} +func New(q *store.Queries, secretHex string, _ string) (*Manager, error) { + m := &Manager{q: q} if secretHex != "" { // Accept both hex (64 chars) and base64 (44 chars) for convenience. var secret []byte @@ -66,16 +65,13 @@ func New(q *store.Queries, secretHex string, fallbackKey string) (*Manager, erro return m, nil } -// Pick selects the best key for a request and returns its decrypted plaintext. -// Records the key id on the returned Selection for later spend recording. +// Pick selects the best key for a request using the v2 picker (health-aware, +// $10 buffer, max_micros cap). Returns ErrNoKeys when the pool has no +// eligible keys — the caller should surface this as a 503 to the user. func (m *Manager) Pick(ctx context.Context) (*Selection, error) { - today := todayUTC() - key, err := m.q.PickPoolKey(ctx, today) + key, err := m.q.PickPoolKeyV2(ctx) if err != nil { if errors.Is(err, sql.ErrNoRows) { - if m.fallbackKey != "" { - return &Selection{keyID: "", plaintext: m.fallbackKey, fallback: true}, nil - } return nil, ErrNoKeys } return nil, fmt.Errorf("pool: pick key: %w", err) @@ -91,12 +87,12 @@ func (m *Manager) Pick(ctx context.Context) (*Selection, error) { // Call this after the upstream call settles (use context.Background() — the // request context may already be canceled). func (m *Manager) RecordSpend(ctx context.Context, sel *Selection, amountMicros int64) error { - if sel.fallback || sel.keyID == "" { - return nil // static fallback key — no pool row to update + if sel.keyID == "" { + return nil } now := time.Now().Unix() return m.q.RecordPoolKeySpend(ctx, store.RecordPoolKeySpendParams{ - TodayDate: todayUTC(), + TodayDate: time.Now().UTC().Unix() / 86400, TodayMicros: amountMicros, LastUsedAt: sql.NullInt64{Int64: now, Valid: true}, ID: sel.keyID, @@ -169,19 +165,10 @@ func Fingerprint(plaintext string) string { type Selection struct { keyID string plaintext string - fallback bool } // APIKey returns the decrypted pioneer API key for use in HTTP headers. func (s *Selection) APIKey() string { return s.plaintext } -// IsFallback reports whether the fallback static key was used. -func (s *Selection) IsFallback() bool { return s.fallback } - -// KeyID returns the DB row ID of the selected key (empty for fallback). +// KeyID returns the DB row ID of the selected key. func (s *Selection) KeyID() string { return s.keyID } - -// todayUTC returns the current UTC day as unix seconds / 86400. -func todayUTC() int64 { - return time.Now().UTC().Unix() / 86400 -} diff --git a/server/internal/pool/reconciler.go b/server/internal/pool/reconciler.go new file mode 100644 index 0000000..9ad5cb1 --- /dev/null +++ b/server/internal/pool/reconciler.go @@ -0,0 +1,324 @@ +// Package pool manages the shared pioneer.ai API key pool. +// +// The reconciler runs on a background ticker and keeps pool_keys in sync +// with pioneer's billing API. See design/pool-keys.md for the full design. +// +// Health integer enum stored in pool_keys.pioneer_health: +// +// 0 = unknown (just added, not yet probed) +// 1 = healthy (last probe succeeded) +// 2 = unauthorized (got 401 — exhausted or revoked, can't tell) +// +// The reconciler marks keys unauthorized on 401 and revokes them after +// UnhealthyRevokeAfter consecutive days. 503 is transient and leaves +// health unchanged. +package pool + +import ( + "context" + "database/sql" + "encoding/json" + "fmt" + "io" + "net/http" + "time" + + charmlog "charm.land/log/v2" + + "github.com/taciturnaxolotl/potluck/internal/store" +) + +const ( + // HealthUnknown is the initial state for a new key. + HealthUnknown = int64(0) + // HealthHealthy means the last pioneer probe succeeded. + HealthHealthy = int64(1) + // HealthUnauthorized means pioneer returned 401. Could be exhausted or revoked. + HealthUnauthorized = int64(2) + + // PoolKeyBufferCredits is the minimum pioneer credits remaining before + // we consider a key exhausted for picking purposes. + // 1000 credits = $10 USD = 10_000_000 potluck micros. + PoolKeyBufferCredits = int64(1000) + PoolKeyBufferMicros = PoolKeyBufferCredits * 10_000 + + // UnhealthyRevokeAfterDays: after this many consecutive days of 401, + // permanently revoke the key. + UnhealthyRevokeAfterDays = 14 + + // ReconcileInterval is how often the reconciler wakes. + ReconcileInterval = 10 * time.Minute + + // Only accept pro-plan pioneer keys. + RequiredPaymentPlan = "pro" +) + +// PlanInfo is the subset of /billing/plan-info we care about. +type PlanInfo struct { + PaymentPlan string `json:"payment_plan"` + CreditLimit float64 `json:"credit_limit"` + TotalUsage float64 `json:"total_usage"` + RemainingCredits float64 `json:"remaining_credits"` +} + +// BillingStatus is the subset of /billing/billing-status we care about. +type BillingStatus struct { + TeamID string `json:"team_id"` +} + +// probeResult holds the outcome of a single-key billing probe. +type probeResult struct { + // err is non-nil for network/decode errors only. + // http401 and http503 are separate flags. + err error + http401 bool + http503 bool + + plan PlanInfo + teamID string + + // todayMicros = total_usage * 10_000 (credits → micros). + todayMicros int64 + // remainingMicros = remaining_credits * 10_000. + remainingMicros int64 + // creditLimitMicros = credit_limit * 10_000. + creditLimitMicros int64 +} + +// probePlanInfo calls /billing/plan-info and /billing/billing-status for a key. +func probePlanInfo(ctx context.Context, httpClient *http.Client, apiKey string) probeResult { + ctx, cancel := context.WithTimeout(ctx, 12*time.Second) + defer cancel() + + do := func(url string) ([]byte, int, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, 0, err + } + req.Header.Set("Authorization", "Bearer "+apiKey) + resp, err := httpClient.Do(req) + if err != nil { + return nil, 0, err + } + defer resp.Body.Close() + b, _ := io.ReadAll(resp.Body) + return b, resp.StatusCode, nil + } + + // /billing/plan-info + body, status, err := do("https://api.pioneer.ai/billing/plan-info") + if err != nil { + return probeResult{err: fmt.Errorf("plan-info: %w", err)} + } + switch status { + case http.StatusUnauthorized: + return probeResult{http401: true} + case http.StatusServiceUnavailable: + return probeResult{http503: true} + } + if status/100 != 2 { + return probeResult{err: fmt.Errorf("plan-info: unexpected status %d", status)} + } + var plan PlanInfo + if err := json.Unmarshal(body, &plan); err != nil { + return probeResult{err: fmt.Errorf("plan-info decode: %w", err)} + } + + // /billing/billing-status (for team_id) + body2, status2, err := do("https://api.pioneer.ai/billing/billing-status") + var teamID string + if err == nil && status2 == http.StatusOK { + var bs BillingStatus + if json.Unmarshal(body2, &bs) == nil { + teamID = bs.TeamID + } + } + + return probeResult{ + plan: plan, + teamID: teamID, + todayMicros: int64(plan.TotalUsage * 10_000), + remainingMicros: int64(plan.RemainingCredits * 10_000), + creditLimitMicros: int64(plan.CreditLimit * 10_000), + } +} + +// Reconciler syncs pool_keys with pioneer's billing API on a ticker. +type Reconciler struct { + q *store.Queries + decrypt func(string) (string, error) + httpClient *http.Client + log *charmlog.Logger +} + +// NewReconciler creates a Reconciler. decrypt is pool.Manager.Decrypt. +func NewReconciler(q *store.Queries, decrypt func(string) (string, error), log *charmlog.Logger) *Reconciler { + return &Reconciler{ + q: q, + decrypt: decrypt, + httpClient: &http.Client{Timeout: 20 * time.Second}, + log: log, + } +} + +// Run starts the reconciler loop. Call in a goroutine; stops when ctx is done. +func (r *Reconciler) Run(ctx context.Context) { + r.log.Info("pool reconciler starting", "interval", ReconcileInterval) + // Run once immediately so we don't wait 10 min on startup. + r.tick(ctx) + + t := time.NewTicker(ReconcileInterval) + defer t.Stop() + for { + select { + case <-ctx.Done(): + r.log.Info("pool reconciler stopping") + return + case <-t.C: + r.tick(ctx) + } + } +} + +// tick runs one reconciliation pass over all non-revoked keys. +func (r *Reconciler) tick(ctx context.Context) { + keys, err := r.q.ListKeysNeedingHealthCheck(ctx) + if err != nil { + r.log.Error("reconciler: list keys", "err", err) + return + } + + // Check for keys to permanently revoke (>14 days of 401). + cutoff := time.Now().Add(-UnhealthyRevokeAfterDays * 24 * time.Hour).Unix() + stale, err := r.q.ListUnhealthyKeysOlderThan(ctx, sql.NullInt64{Int64: cutoff, Valid: true}) + if err != nil { + r.log.Error("reconciler: list unhealthy", "err", err) + } + for _, k := range stale { + r.log.Warn("pool reconciler: revoking key after 14 days of 401", + "key_id", k.ID, "label", k.Label, + "unhealthy_since", k.PioneerUnhealthySince.Int64) + _ = r.q.MarkPoolKeyRevoked(ctx, store.MarkPoolKeyRevokedParams{ + RevokedAt: sql.NullInt64{Int64: time.Now().Unix(), Valid: true}, + ID: k.ID, + }) + } + + for _, key := range keys { + // Skip if already permanently revoked (race with above). + if key.RevokedAt.Valid { + continue + } + r.probeKey(ctx, key) + } +} + +// probeKey probes a single key and updates its DB row. +func (r *Reconciler) probeKey(ctx context.Context, key store.PoolKey) { + plaintext, err := r.decrypt(key.KeyCiphertext) + if err != nil { + r.log.Error("reconciler: decrypt key", "key_id", key.ID, "err", err) + return + } + + result := probePlanInfo(ctx, r.httpClient, plaintext) + now := time.Now().Unix() + + switch { + case result.http503: + // Pioneer auth service down — transient, don't touch health. + r.log.Warn("reconciler: pioneer auth down (503), skipping key", + "key_id", key.ID, "label", key.Label) + return + + case result.http401: + // Key exhausted or revoked — mark unauthorized. + var unhealthySince sql.NullInt64 + if key.PioneerHealth == HealthUnauthorized && key.PioneerUnhealthySince.Valid { + // Preserve existing timestamp. + unhealthySince = key.PioneerUnhealthySince + } else { + unhealthySince = sql.NullInt64{Int64: now, Valid: true} + } + r.log.Info("reconciler: key unauthorized (401)", + "key_id", key.ID, "label", key.Label, + "unhealthy_since", unhealthySince.Int64) + _ = r.q.UpdatePoolKeyHealth(ctx, store.UpdatePoolKeyHealthParams{ + PioneerHealth: HealthUnauthorized, + PioneerUnhealthySince: unhealthySince, + PioneerTeamID: key.PioneerTeamID, + PioneerPaymentPlan: key.PioneerPaymentPlan, + PioneerCreditLimitMicros: key.PioneerCreditLimitMicros, + PioneerRemainingMicros: key.PioneerRemainingMicros, + TodayMicros: key.TodayMicros, + LastBillingSyncAt: sql.NullInt64{Int64: now, Valid: true}, + ID: key.ID, + }) + return + + case result.err != nil: + // Network/decode error — transient, log and skip. + r.log.Error("reconciler: probe failed", + "key_id", key.ID, "label", key.Label, "err", result.err) + return + } + + // Successful probe — validate plan and update. + if result.plan.PaymentPlan != RequiredPaymentPlan { + r.log.Warn("reconciler: non-pro key, marking unauthorized", + "key_id", key.ID, "label", key.Label, + "plan", result.plan.PaymentPlan) + _ = r.q.UpdatePoolKeyHealth(ctx, store.UpdatePoolKeyHealthParams{ + PioneerHealth: HealthUnauthorized, + PioneerUnhealthySince: sql.NullInt64{Int64: now, Valid: true}, + PioneerTeamID: nullStr(result.teamID), + PioneerPaymentPlan: nullStr(result.plan.PaymentPlan), + PioneerCreditLimitMicros: nullInt(result.creditLimitMicros), + PioneerRemainingMicros: nullInt(result.remainingMicros), + TodayMicros: result.todayMicros, + LastBillingSyncAt: sql.NullInt64{Int64: now, Valid: true}, + ID: key.ID, + }) + return + } + + // Was previously unhealthy — reactivate. + if key.PioneerHealth != HealthHealthy || key.PendingValidation == 1 { + r.log.Info("reconciler: key recovered, activating", + "key_id", key.ID, "label", key.Label) + _ = r.q.ActivatePoolKey(ctx, key.ID) + } + + _ = r.q.UpdatePoolKeyHealth(ctx, store.UpdatePoolKeyHealthParams{ + PioneerHealth: HealthHealthy, + PioneerUnhealthySince: sql.NullInt64{}, // NULL — clear it + PioneerTeamID: nullStr(result.teamID), + PioneerPaymentPlan: nullStr(result.plan.PaymentPlan), + PioneerCreditLimitMicros: nullInt(result.creditLimitMicros), + PioneerRemainingMicros: nullInt(result.remainingMicros), + TodayMicros: result.todayMicros, + LastBillingSyncAt: sql.NullInt64{Int64: now, Valid: true}, + ID: key.ID, + }) + + r.log.Info("reconciler: key synced", + "key_id", key.ID, + "label", key.Label, + "today_usd", fmt.Sprintf("%.2f", float64(result.todayMicros)/1_000_000), + "remaining_usd", fmt.Sprintf("%.2f", float64(result.remainingMicros)/1_000_000), + ) +} + +func nullStr(s string) sql.NullString { + if s == "" { + return sql.NullString{} + } + return sql.NullString{String: s, Valid: true} +} + +func nullInt(v int64) sql.NullInt64 { + if v == 0 { + return sql.NullInt64{} + } + return sql.NullInt64{Int64: v, Valid: true} +} diff --git a/server/internal/store/billing_rows.sql.go b/server/internal/store/billing_rows.sql.go new file mode 100644 index 0000000..513013a --- /dev/null +++ b/server/internal/store/billing_rows.sql.go @@ -0,0 +1,183 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 +// source: billing_rows.sql + +package store + +import ( + "context" + "database/sql" +) + +const latestBillingRowTime = `-- name: LatestBillingRowTime :one +SELECT COALESCE(SUM(pioneer_created_at * 0) + MAX(pioneer_created_at), 0) +FROM pool_key_billing_rows +WHERE pool_key_id = ? +` + +// Most recent pioneer_created_at for a key. +// COALESCE(SUM(pioneer_created_at*0) + MAX(...)) is a workaround for +// sqlc's sqlite parser rejecting bare COALESCE(MAX(...), 0). +func (q *Queries) LatestBillingRowTime(ctx context.Context, poolKeyID string) (interface{}, error) { + row := q.db.QueryRowContext(ctx, latestBillingRowTime, poolKeyID) + var coalesce interface{} + err := row.Scan(&coalesce) + return coalesce, err +} + +const listBillingRowsForKeyAfter = `-- name: ListBillingRowsForKeyAfter :many +SELECT id, pool_key_id, pioneer_created_at, credit_micros, cost_micros, token_usage, model, endpoint, attributed_user_id, attribution, is_duplicate, matched_request_id, ingested_at FROM pool_key_billing_rows +WHERE pool_key_id = ? + AND pioneer_created_at > ? +ORDER BY pioneer_created_at ASC +` + +type ListBillingRowsForKeyAfterParams struct { + PoolKeyID string `json:"pool_key_id"` + PioneerCreatedAt int64 `json:"pioneer_created_at"` +} + +// All billing rows for a key after a given timestamp. +func (q *Queries) ListBillingRowsForKeyAfter(ctx context.Context, arg ListBillingRowsForKeyAfterParams) ([]PoolKeyBillingRow, error) { + rows, err := q.db.QueryContext(ctx, listBillingRowsForKeyAfter, arg.PoolKeyID, arg.PioneerCreatedAt) + if err != nil { + return nil, err + } + defer rows.Close() + items := []PoolKeyBillingRow{} + for rows.Next() { + var i PoolKeyBillingRow + if err := rows.Scan( + &i.ID, + &i.PoolKeyID, + &i.PioneerCreatedAt, + &i.CreditMicros, + &i.CostMicros, + &i.TokenUsage, + &i.Model, + &i.Endpoint, + &i.AttributedUserID, + &i.Attribution, + &i.IsDuplicate, + &i.MatchedRequestID, + &i.IngestedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const listBillingRowsForUserToday = `-- name: ListBillingRowsForUserToday :many +SELECT id, pool_key_id, pioneer_created_at, credit_micros, cost_micros, token_usage, model, endpoint, attributed_user_id, attribution, is_duplicate, matched_request_id, ingested_at FROM pool_key_billing_rows +WHERE attributed_user_id = ? + AND pioneer_created_at >= ? + AND pioneer_created_at < ? +ORDER BY pioneer_created_at ASC +` + +type ListBillingRowsForUserTodayParams struct { + AttributedUserID sql.NullString `json:"attributed_user_id"` + PioneerCreatedAt int64 `json:"pioneer_created_at"` + PioneerCreatedAt_2 int64 `json:"pioneer_created_at_2"` +} + +// Billing rows attributed to a user in a time window. +// Caller filters duplicates and sums in Go. +func (q *Queries) ListBillingRowsForUserToday(ctx context.Context, arg ListBillingRowsForUserTodayParams) ([]PoolKeyBillingRow, error) { + rows, err := q.db.QueryContext(ctx, listBillingRowsForUserToday, arg.AttributedUserID, arg.PioneerCreatedAt, arg.PioneerCreatedAt_2) + if err != nil { + return nil, err + } + defer rows.Close() + items := []PoolKeyBillingRow{} + for rows.Next() { + var i PoolKeyBillingRow + if err := rows.Scan( + &i.ID, + &i.PoolKeyID, + &i.PioneerCreatedAt, + &i.CreditMicros, + &i.CostMicros, + &i.TokenUsage, + &i.Model, + &i.Endpoint, + &i.AttributedUserID, + &i.Attribution, + &i.IsDuplicate, + &i.MatchedRequestID, + &i.IngestedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const upsertBillingRow = `-- name: UpsertBillingRow :exec + +INSERT INTO pool_key_billing_rows ( + id, pool_key_id, pioneer_created_at, credit_micros, cost_micros, + token_usage, model, endpoint, + attributed_user_id, attribution, is_duplicate, matched_request_id, ingested_at +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +ON CONFLICT(id) DO NOTHING +` + +type UpsertBillingRowParams struct { + ID string `json:"id"` + PoolKeyID string `json:"pool_key_id"` + PioneerCreatedAt int64 `json:"pioneer_created_at"` + CreditMicros int64 `json:"credit_micros"` + CostMicros int64 `json:"cost_micros"` + TokenUsage int64 `json:"token_usage"` + Model string `json:"model"` + Endpoint string `json:"endpoint"` + AttributedUserID sql.NullString `json:"attributed_user_id"` + Attribution int64 `json:"attribution"` + IsDuplicate int64 `json:"is_duplicate"` + MatchedRequestID sql.NullString `json:"matched_request_id"` + IngestedAt int64 `json:"ingested_at"` +} + +// pool_key_billing_rows: ingested pioneer billing log entries. +// +// attribution integer enum: +// +// 0 = matched (joined to a potluck_requests row) +// 1 = judge_paired (/llmaj/judge paired to preceding opus call) +// 2 = owner_fallback (no match, charged to key owner) +// 3 = duplicate (double-logged by pioneer, not charged) +func (q *Queries) UpsertBillingRow(ctx context.Context, arg UpsertBillingRowParams) error { + _, err := q.db.ExecContext(ctx, upsertBillingRow, + arg.ID, + arg.PoolKeyID, + arg.PioneerCreatedAt, + arg.CreditMicros, + arg.CostMicros, + arg.TokenUsage, + arg.Model, + arg.Endpoint, + arg.AttributedUserID, + arg.Attribution, + arg.IsDuplicate, + arg.MatchedRequestID, + arg.IngestedAt, + ) + return err +} diff --git a/server/internal/store/daily_spend.sql.go b/server/internal/store/daily_spend.sql.go new file mode 100644 index 0000000..fd3257b --- /dev/null +++ b/server/internal/store/daily_spend.sql.go @@ -0,0 +1,197 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 +// source: daily_spend.sql + +package store + +import ( + "context" +) + +const getLatestRecompute = `-- name: GetLatestRecompute :one +SELECT set_at, set_by_user_id FROM user_daily_allowances +WHERE day = ? +ORDER BY set_at DESC +LIMIT 1 +` + +type GetLatestRecomputeRow struct { + SetAt int64 `json:"set_at"` + SetByUserID string `json:"set_by_user_id"` +} + +// Most recent recompute for today, for display in the dashboard. +func (q *Queries) GetLatestRecompute(ctx context.Context, day int64) (GetLatestRecomputeRow, error) { + row := q.db.QueryRowContext(ctx, getLatestRecompute, day) + var i GetLatestRecomputeRow + err := row.Scan(&i.SetAt, &i.SetByUserID) + return i, err +} + +const getUserDailyAllowance = `-- name: GetUserDailyAllowance :one +SELECT user_id, day, shared_allowance_micros, set_at, set_by_user_id FROM user_daily_allowances WHERE user_id = ? AND day = ? +` + +type GetUserDailyAllowanceParams struct { + UserID string `json:"user_id"` + Day int64 `json:"day"` +} + +func (q *Queries) GetUserDailyAllowance(ctx context.Context, arg GetUserDailyAllowanceParams) (UserDailyAllowance, error) { + row := q.db.QueryRowContext(ctx, getUserDailyAllowance, arg.UserID, arg.Day) + var i UserDailyAllowance + err := row.Scan( + &i.UserID, + &i.Day, + &i.SharedAllowanceMicros, + &i.SetAt, + &i.SetByUserID, + ) + return i, err +} + +const getUserDailySpend = `-- name: GetUserDailySpend :one +SELECT user_id, day, shared_spent_micros, private_spent_micros FROM user_daily_spend WHERE user_id = ? AND day = ? +` + +type GetUserDailySpendParams struct { + UserID string `json:"user_id"` + Day int64 `json:"day"` +} + +func (q *Queries) GetUserDailySpend(ctx context.Context, arg GetUserDailySpendParams) (UserDailySpend, error) { + row := q.db.QueryRowContext(ctx, getUserDailySpend, arg.UserID, arg.Day) + var i UserDailySpend + err := row.Scan( + &i.UserID, + &i.Day, + &i.SharedSpentMicros, + &i.PrivateSpentMicros, + ) + return i, err +} + +const listUserDailyAllowancesForDay = `-- name: ListUserDailyAllowancesForDay :many +SELECT user_id, day, shared_allowance_micros, set_at, set_by_user_id FROM user_daily_allowances WHERE day = ? +` + +func (q *Queries) ListUserDailyAllowancesForDay(ctx context.Context, day int64) ([]UserDailyAllowance, error) { + rows, err := q.db.QueryContext(ctx, listUserDailyAllowancesForDay, day) + if err != nil { + return nil, err + } + defer rows.Close() + items := []UserDailyAllowance{} + for rows.Next() { + var i UserDailyAllowance + if err := rows.Scan( + &i.UserID, + &i.Day, + &i.SharedAllowanceMicros, + &i.SetAt, + &i.SetByUserID, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const listUserDailySpendForDay = `-- name: ListUserDailySpendForDay :many +SELECT user_id, day, shared_spent_micros, private_spent_micros FROM user_daily_spend WHERE day = ? +` + +func (q *Queries) ListUserDailySpendForDay(ctx context.Context, day int64) ([]UserDailySpend, error) { + rows, err := q.db.QueryContext(ctx, listUserDailySpendForDay, day) + if err != nil { + return nil, err + } + defer rows.Close() + items := []UserDailySpend{} + for rows.Next() { + var i UserDailySpend + if err := rows.Scan( + &i.UserID, + &i.Day, + &i.SharedSpentMicros, + &i.PrivateSpentMicros, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const upsertUserDailyAllowance = `-- name: UpsertUserDailyAllowance :exec +INSERT INTO user_daily_allowances ( + user_id, day, shared_allowance_micros, set_at, set_by_user_id +) VALUES (?, ?, ?, ?, ?) +ON CONFLICT(user_id, day) DO UPDATE SET + shared_allowance_micros = MAX( + excluded.shared_allowance_micros, + user_daily_allowances.shared_allowance_micros + ), + set_at = excluded.set_at, + set_by_user_id = excluded.set_by_user_id +` + +type UpsertUserDailyAllowanceParams struct { + UserID string `json:"user_id"` + Day int64 `json:"day"` + SharedAllowanceMicros int64 `json:"shared_allowance_micros"` + SetAt int64 `json:"set_at"` + SetByUserID string `json:"set_by_user_id"` +} + +func (q *Queries) UpsertUserDailyAllowance(ctx context.Context, arg UpsertUserDailyAllowanceParams) error { + _, err := q.db.ExecContext(ctx, upsertUserDailyAllowance, + arg.UserID, + arg.Day, + arg.SharedAllowanceMicros, + arg.SetAt, + arg.SetByUserID, + ) + return err +} + +const upsertUserDailySpend = `-- name: UpsertUserDailySpend :exec + +INSERT INTO user_daily_spend (user_id, day, shared_spent_micros, private_spent_micros) +VALUES (?, ?, ?, ?) +ON CONFLICT(user_id, day) DO UPDATE SET + shared_spent_micros = shared_spent_micros + excluded.shared_spent_micros, + private_spent_micros = private_spent_micros + excluded.private_spent_micros +` + +type UpsertUserDailySpendParams struct { + UserID string `json:"user_id"` + Day int64 `json:"day"` + SharedSpentMicros int64 `json:"shared_spent_micros"` + PrivateSpentMicros int64 `json:"private_spent_micros"` +} + +// user_daily_spend and user_daily_allowances. +func (q *Queries) UpsertUserDailySpend(ctx context.Context, arg UpsertUserDailySpendParams) error { + _, err := q.db.ExecContext(ctx, upsertUserDailySpend, + arg.UserID, + arg.Day, + arg.SharedSpentMicros, + arg.PrivateSpentMicros, + ) + return err +} diff --git a/server/internal/store/models.go b/server/internal/store/models.go index 5b3a2a8..5305f22 100644 --- a/server/internal/store/models.go +++ b/server/internal/store/models.go @@ -68,20 +68,77 @@ type ModelPrice struct { UpdatedAt int64 `json:"updated_at"` } +type ModelsCatalog struct { + ID string `json:"id"` + Label string `json:"label"` + Description string `json:"description"` + ContextWindow sql.NullInt64 `json:"context_window"` + MaxOutputTokens sql.NullInt64 `json:"max_output_tokens"` + IsChat int64 `json:"is_chat"` + Tier sql.NullString `json:"tier"` + InputPricePerMillionMicros sql.NullInt64 `json:"input_price_per_million_micros"` + OutputPricePerMillionMicros sql.NullInt64 `json:"output_price_per_million_micros"` + RawJson string `json:"raw_json"` + RefreshedAt int64 `json:"refreshed_at"` +} + type PoolKey struct { - ID string `json:"id"` - UserID string `json:"user_id"` - Label string `json:"label"` - KeyCiphertext string `json:"key_ciphertext"` - KeyFingerprint string `json:"key_fingerprint"` - Active int64 `json:"active"` - DailyLimitMicros int64 `json:"daily_limit_micros"` - TodayDate int64 `json:"today_date"` - TodayMicros int64 `json:"today_micros"` - TotalMicros int64 `json:"total_micros"` - RequestCount int64 `json:"request_count"` - CreatedAt int64 `json:"created_at"` - LastUsedAt sql.NullInt64 `json:"last_used_at"` + ID string `json:"id"` + UserID string `json:"user_id"` + Label string `json:"label"` + KeyCiphertext string `json:"key_ciphertext"` + KeyFingerprint string `json:"key_fingerprint"` + Active int64 `json:"active"` + DailyLimitMicros int64 `json:"daily_limit_micros"` + TodayDate int64 `json:"today_date"` + TodayMicros int64 `json:"today_micros"` + TotalMicros int64 `json:"total_micros"` + RequestCount int64 `json:"request_count"` + CreatedAt int64 `json:"created_at"` + LastUsedAt sql.NullInt64 `json:"last_used_at"` + MaxMicros int64 `json:"max_micros"` + SharedMicros int64 `json:"shared_micros"` + PioneerTeamID sql.NullString `json:"pioneer_team_id"` + PioneerPaymentPlan sql.NullString `json:"pioneer_payment_plan"` + PioneerCreditLimitMicros sql.NullInt64 `json:"pioneer_credit_limit_micros"` + PioneerRemainingMicros sql.NullInt64 `json:"pioneer_remaining_micros"` + PioneerHealth int64 `json:"pioneer_health"` + PioneerUnhealthySince sql.NullInt64 `json:"pioneer_unhealthy_since"` + PendingValidation int64 `json:"pending_validation"` + LastBillingSyncAt sql.NullInt64 `json:"last_billing_sync_at"` + RevokedAt sql.NullInt64 `json:"revoked_at"` +} + +type PoolKeyBillingRow struct { + ID string `json:"id"` + PoolKeyID string `json:"pool_key_id"` + PioneerCreatedAt int64 `json:"pioneer_created_at"` + CreditMicros int64 `json:"credit_micros"` + CostMicros int64 `json:"cost_micros"` + TokenUsage int64 `json:"token_usage"` + Model string `json:"model"` + Endpoint string `json:"endpoint"` + AttributedUserID sql.NullString `json:"attributed_user_id"` + Attribution int64 `json:"attribution"` + IsDuplicate int64 `json:"is_duplicate"` + MatchedRequestID sql.NullString `json:"matched_request_id"` + IngestedAt int64 `json:"ingested_at"` +} + +type PotluckRequest struct { + ID string `json:"id"` + UserID string `json:"user_id"` + ApiKeyID sql.NullString `json:"api_key_id"` + PoolKeyID sql.NullString `json:"pool_key_id"` + Surface string `json:"surface"` + Model string `json:"model"` + StartedAt int64 `json:"started_at"` + FinishedAt sql.NullInt64 `json:"finished_at"` + PromptTokens sql.NullInt64 `json:"prompt_tokens"` + CompletionTokens sql.NullInt64 `json:"completion_tokens"` + TotalTokens sql.NullInt64 `json:"total_tokens"` + Status string `json:"status"` + ErrorCode sql.NullString `json:"error_code"` } type Session struct { @@ -139,3 +196,18 @@ type User struct { SlackID sql.NullString `json:"slack_id"` VerificationStatus sql.NullString `json:"verification_status"` } + +type UserDailyAllowance struct { + UserID string `json:"user_id"` + Day int64 `json:"day"` + SharedAllowanceMicros int64 `json:"shared_allowance_micros"` + SetAt int64 `json:"set_at"` + SetByUserID string `json:"set_by_user_id"` +} + +type UserDailySpend struct { + UserID string `json:"user_id"` + Day int64 `json:"day"` + SharedSpentMicros int64 `json:"shared_spent_micros"` + PrivateSpentMicros int64 `json:"private_spent_micros"` +} diff --git a/server/internal/store/models_catalog.sql.go b/server/internal/store/models_catalog.sql.go new file mode 100644 index 0000000..76da2d6 --- /dev/null +++ b/server/internal/store/models_catalog.sql.go @@ -0,0 +1,117 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 +// source: models_catalog.sql + +package store + +import ( + "context" + "database/sql" +) + +const getModelCatalogRefreshedAt = `-- name: GetModelCatalogRefreshedAt :one +SELECT COALESCE(MIN(refreshed_at), 0) FROM models_catal +` + +// Oldest refreshed_at across all models — tells us when the catalog is stale. +func (q *Queries) GetModelCatalogRefreshedAt(ctx context.Context) (interface{}, error) { + row := q.db.QueryRowContext(ctx, getModelCatalogRefreshedAt) + var coalesce interface{} + err := row.Scan(&coalesce) + return coalesce, err +} + +const listModelCatalog = `-- name: ListModelCatalog :many +SELECT id, label, description, context_window, max_output_tokens, is_chat, tier, input_price_per_million_micros, output_price_per_million_micros, raw_json, refreshed_at FROM models_catalog +WHERE is_chat = 1 +ORDER BY tier ASC, label ASC +` + +func (q *Queries) ListModelCatalog(ctx context.Context) ([]ModelsCatalog, error) { + rows, err := q.db.QueryContext(ctx, listModelCatalog) + if err != nil { + return nil, err + } + defer rows.Close() + items := []ModelsCatalog{} + for rows.Next() { + var i ModelsCatalog + if err := rows.Scan( + &i.ID, + &i.Label, + &i.Description, + &i.ContextWindow, + &i.MaxOutputTokens, + &i.IsChat, + &i.Tier, + &i.InputPricePerMillionMicros, + &i.OutputPricePerMillionMicros, + &i.RawJson, + &i.RefreshedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const upsertModelCatalog = `-- name: UpsertModelCatalog :exec + +INSERT INTO models_catalog ( + id, label, description, context_window, max_output_tokens, + is_chat, tier, + input_price_per_million_micros, output_price_per_million_micros, + raw_json, refreshed_at +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +ON CONFLICT(id) DO UPDATE SET + label = excluded.label, + description = excluded.description, + context_window = excluded.context_window, + max_output_tokens = excluded.max_output_tokens, + is_chat = excluded.is_chat, + tier = excluded.tier, + input_price_per_million_micros = excluded.input_price_per_million_micros, + output_price_per_million_micros = excluded.output_price_per_million_micros, + raw_json = excluded.raw_json, + refreshed_at = excluded.refreshed_at +` + +type UpsertModelCatalogParams struct { + ID string `json:"id"` + Label string `json:"label"` + Description string `json:"description"` + ContextWindow sql.NullInt64 `json:"context_window"` + MaxOutputTokens sql.NullInt64 `json:"max_output_tokens"` + IsChat int64 `json:"is_chat"` + Tier sql.NullString `json:"tier"` + InputPricePerMillionMicros sql.NullInt64 `json:"input_price_per_million_micros"` + OutputPricePerMillionMicros sql.NullInt64 `json:"output_price_per_million_micros"` + RawJson string `json:"raw_json"` + RefreshedAt int64 `json:"refreshed_at"` +} + +// models_catalog: hourly-refreshed cache of /v1/models + /base-models. +func (q *Queries) UpsertModelCatalog(ctx context.Context, arg UpsertModelCatalogParams) error { + _, err := q.db.ExecContext(ctx, upsertModelCatalog, + arg.ID, + arg.Label, + arg.Description, + arg.ContextWindow, + arg.MaxOutputTokens, + arg.IsChat, + arg.Tier, + arg.InputPricePerMillionMicros, + arg.OutputPricePerMillionMicros, + arg.RawJson, + arg.RefreshedAt, + ) + return err +} diff --git a/server/internal/store/pool_keys.sql.go b/server/internal/store/pool_keys.sql.go index ff14d00..e6d21f6 100644 --- a/server/internal/store/pool_keys.sql.go +++ b/server/internal/store/pool_keys.sql.go @@ -17,7 +17,7 @@ INSERT INTO pool_keys ( active, daily_limit_micros, today_date, today_micros, total_micros, request_count, created_at ) VALUES (?, ?, ?, ?, ?, 1, ?, 0, 0, 0, 0, ?) -RETURNING id, user_id, label, key_ciphertext, key_fingerprint, active, daily_limit_micros, today_date, today_micros, total_micros, request_count, created_at, last_used_at +RETURNING id, user_id, label, key_ciphertext, key_fingerprint, active, daily_limit_micros, today_date, today_micros, total_micros, request_count, created_at, last_used_at, max_micros, shared_micros, pioneer_team_id, pioneer_payment_plan, pioneer_credit_limit_micros, pioneer_remaining_micros, pioneer_health, pioneer_unhealthy_since, pending_validation, last_billing_sync_at, revoked_at ` type CreatePoolKeyParams struct { @@ -61,6 +61,17 @@ func (q *Queries) CreatePoolKey(ctx context.Context, arg CreatePoolKeyParams) (P &i.RequestCount, &i.CreatedAt, &i.LastUsedAt, + &i.MaxMicros, + &i.SharedMicros, + &i.PioneerTeamID, + &i.PioneerPaymentPlan, + &i.PioneerCreditLimitMicros, + &i.PioneerRemainingMicros, + &i.PioneerHealth, + &i.PioneerUnhealthySince, + &i.PendingValidation, + &i.LastBillingSyncAt, + &i.RevokedAt, ) return i, err } @@ -80,7 +91,7 @@ func (q *Queries) DeletePoolKey(ctx context.Context, arg DeletePoolKeyParams) er } const getPoolKey = `-- name: GetPoolKey :one -SELECT id, user_id, label, key_ciphertext, key_fingerprint, active, daily_limit_micros, today_date, today_micros, total_micros, request_count, created_at, last_used_at FROM pool_keys WHERE id = ? +SELECT id, user_id, label, key_ciphertext, key_fingerprint, active, daily_limit_micros, today_date, today_micros, total_micros, request_count, created_at, last_used_at, max_micros, shared_micros, pioneer_team_id, pioneer_payment_plan, pioneer_credit_limit_micros, pioneer_remaining_micros, pioneer_health, pioneer_unhealthy_since, pending_validation, last_billing_sync_at, revoked_at FROM pool_keys WHERE id = ? ` func (q *Queries) GetPoolKey(ctx context.Context, id string) (PoolKey, error) { @@ -100,6 +111,17 @@ func (q *Queries) GetPoolKey(ctx context.Context, id string) (PoolKey, error) { &i.RequestCount, &i.CreatedAt, &i.LastUsedAt, + &i.MaxMicros, + &i.SharedMicros, + &i.PioneerTeamID, + &i.PioneerPaymentPlan, + &i.PioneerCreditLimitMicros, + &i.PioneerRemainingMicros, + &i.PioneerHealth, + &i.PioneerUnhealthySince, + &i.PendingValidation, + &i.LastBillingSyncAt, + &i.RevokedAt, ) return i, err } @@ -166,28 +188,39 @@ func (q *Queries) ListPoolAllocations(ctx context.Context) ([]ListPoolAllocation } const listPoolKeys = `-- name: ListPoolKeys :many -SELECT pk.id, pk.user_id, pk.label, pk.key_ciphertext, pk.key_fingerprint, pk.active, pk.daily_limit_micros, pk.today_date, pk.today_micros, pk.total_micros, pk.request_count, pk.created_at, pk.last_used_at, u.display_name AS owner_name, u.email AS owner_email +SELECT pk.id, pk.user_id, pk.label, pk.key_ciphertext, pk.key_fingerprint, pk.active, pk.daily_limit_micros, pk.today_date, pk.today_micros, pk.total_micros, pk.request_count, pk.created_at, pk.last_used_at, pk.max_micros, pk.shared_micros, pk.pioneer_team_id, pk.pioneer_payment_plan, pk.pioneer_credit_limit_micros, pk.pioneer_remaining_micros, pk.pioneer_health, pk.pioneer_unhealthy_since, pk.pending_validation, pk.last_billing_sync_at, pk.revoked_at, u.display_name AS owner_name, u.email AS owner_email FROM pool_keys pk JOIN users u ON u.id = pk.user_id ORDER BY pk.active DESC, pk.today_micros ASC ` type ListPoolKeysRow struct { - ID string `json:"id"` - UserID string `json:"user_id"` - Label string `json:"label"` - KeyCiphertext string `json:"key_ciphertext"` - KeyFingerprint string `json:"key_fingerprint"` - Active int64 `json:"active"` - DailyLimitMicros int64 `json:"daily_limit_micros"` - TodayDate int64 `json:"today_date"` - TodayMicros int64 `json:"today_micros"` - TotalMicros int64 `json:"total_micros"` - RequestCount int64 `json:"request_count"` - CreatedAt int64 `json:"created_at"` - LastUsedAt sql.NullInt64 `json:"last_used_at"` - OwnerName string `json:"owner_name"` - OwnerEmail string `json:"owner_email"` + ID string `json:"id"` + UserID string `json:"user_id"` + Label string `json:"label"` + KeyCiphertext string `json:"key_ciphertext"` + KeyFingerprint string `json:"key_fingerprint"` + Active int64 `json:"active"` + DailyLimitMicros int64 `json:"daily_limit_micros"` + TodayDate int64 `json:"today_date"` + TodayMicros int64 `json:"today_micros"` + TotalMicros int64 `json:"total_micros"` + RequestCount int64 `json:"request_count"` + CreatedAt int64 `json:"created_at"` + LastUsedAt sql.NullInt64 `json:"last_used_at"` + MaxMicros int64 `json:"max_micros"` + SharedMicros int64 `json:"shared_micros"` + PioneerTeamID sql.NullString `json:"pioneer_team_id"` + PioneerPaymentPlan sql.NullString `json:"pioneer_payment_plan"` + PioneerCreditLimitMicros sql.NullInt64 `json:"pioneer_credit_limit_micros"` + PioneerRemainingMicros sql.NullInt64 `json:"pioneer_remaining_micros"` + PioneerHealth int64 `json:"pioneer_health"` + PioneerUnhealthySince sql.NullInt64 `json:"pioneer_unhealthy_since"` + PendingValidation int64 `json:"pending_validation"` + LastBillingSyncAt sql.NullInt64 `json:"last_billing_sync_at"` + RevokedAt sql.NullInt64 `json:"revoked_at"` + OwnerName string `json:"owner_name"` + OwnerEmail string `json:"owner_email"` } // All keys (for the pool management page). Includes inactive and other users' keys. @@ -214,6 +247,17 @@ func (q *Queries) ListPoolKeys(ctx context.Context) ([]ListPoolKeysRow, error) { &i.RequestCount, &i.CreatedAt, &i.LastUsedAt, + &i.MaxMicros, + &i.SharedMicros, + &i.PioneerTeamID, + &i.PioneerPaymentPlan, + &i.PioneerCreditLimitMicros, + &i.PioneerRemainingMicros, + &i.PioneerHealth, + &i.PioneerUnhealthySince, + &i.PendingValidation, + &i.LastBillingSyncAt, + &i.RevokedAt, &i.OwnerName, &i.OwnerEmail, ); err != nil { @@ -231,7 +275,7 @@ func (q *Queries) ListPoolKeys(ctx context.Context) ([]ListPoolKeysRow, error) { } const listPoolKeysForUser = `-- name: ListPoolKeysForUser :many -SELECT id, user_id, label, key_ciphertext, key_fingerprint, active, daily_limit_micros, today_date, today_micros, total_micros, request_count, created_at, last_used_at FROM pool_keys WHERE user_id = ? ORDER BY created_at DESC +SELECT id, user_id, label, key_ciphertext, key_fingerprint, active, daily_limit_micros, today_date, today_micros, total_micros, request_count, created_at, last_used_at, max_micros, shared_micros, pioneer_team_id, pioneer_payment_plan, pioneer_credit_limit_micros, pioneer_remaining_micros, pioneer_health, pioneer_unhealthy_since, pending_validation, last_billing_sync_at, revoked_at FROM pool_keys WHERE user_id = ? ORDER BY created_at DESC ` func (q *Queries) ListPoolKeysForUser(ctx context.Context, userID string) ([]PoolKey, error) { @@ -257,6 +301,17 @@ func (q *Queries) ListPoolKeysForUser(ctx context.Context, userID string) ([]Poo &i.RequestCount, &i.CreatedAt, &i.LastUsedAt, + &i.MaxMicros, + &i.SharedMicros, + &i.PioneerTeamID, + &i.PioneerPaymentPlan, + &i.PioneerCreditLimitMicros, + &i.PioneerRemainingMicros, + &i.PioneerHealth, + &i.PioneerUnhealthySince, + &i.PendingValidation, + &i.LastBillingSyncAt, + &i.RevokedAt, ); err != nil { return nil, err } @@ -272,7 +327,7 @@ func (q *Queries) ListPoolKeysForUser(ctx context.Context, userID string) ([]Poo } const pickPoolKey = `-- name: PickPoolKey :one -SELECT id, user_id, label, key_ciphertext, key_fingerprint, active, daily_limit_micros, today_date, today_micros, total_micros, request_count, created_at, last_used_at FROM pool_keys +SELECT id, user_id, label, key_ciphertext, key_fingerprint, active, daily_limit_micros, today_date, today_micros, total_micros, request_count, created_at, last_used_at, max_micros, shared_micros, pioneer_team_id, pioneer_payment_plan, pioneer_credit_limit_micros, pioneer_remaining_micros, pioneer_health, pioneer_unhealthy_since, pending_validation, last_billing_sync_at, revoked_at FROM pool_keys WHERE active = 1 AND (today_date < ?1 OR today_micros < daily_limit_micros) ORDER BY @@ -302,6 +357,17 @@ func (q *Queries) PickPoolKey(ctx context.Context, todayDate int64) (PoolKey, er &i.RequestCount, &i.CreatedAt, &i.LastUsedAt, + &i.MaxMicros, + &i.SharedMicros, + &i.PioneerTeamID, + &i.PioneerPaymentPlan, + &i.PioneerCreditLimitMicros, + &i.PioneerRemainingMicros, + &i.PioneerHealth, + &i.PioneerUnhealthySince, + &i.PendingValidation, + &i.LastBillingSyncAt, + &i.RevokedAt, ) return i, err } diff --git a/server/internal/store/pool_keys_v2.sql.go b/server/internal/store/pool_keys_v2.sql.go new file mode 100644 index 0000000..67ae80b --- /dev/null +++ b/server/internal/store/pool_keys_v2.sql.go @@ -0,0 +1,286 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 +// source: pool_keys_v2.sql + +package store + +import ( + "context" + "database/sql" +) + +const activatePoolKey = `-- name: ActivatePoolKey :exec +UPDATE pool_keys SET + active = 1, + pending_validation = 0, + pioneer_health = 1, + pioneer_unhealthy_since = NULL +WHERE id = ? +` + +// Called by the reconciler when a pending_validation or unauthorized key +// comes back healthy (pioneer_health=1). +func (q *Queries) ActivatePoolKey(ctx context.Context, id string) error { + _, err := q.db.ExecContext(ctx, activatePoolKey, id) + return err +} + +const listKeysNeedingHealthCheck = `-- name: ListKeysNeedingHealthCheck :many +SELECT id, user_id, label, key_ciphertext, key_fingerprint, active, daily_limit_micros, today_date, today_micros, total_micros, request_count, created_at, last_used_at, max_micros, shared_micros, pioneer_team_id, pioneer_payment_plan, pioneer_credit_limit_micros, pioneer_remaining_micros, pioneer_health, pioneer_unhealthy_since, pending_validation, last_billing_sync_at, revoked_at FROM pool_keys +WHERE revoked_at IS NULL +ORDER BY last_billing_sync_at ASC NULLS FIRST +` + +// Keys the reconciler should probe on each tick. +// Excludes permanently revoked keys. +func (q *Queries) ListKeysNeedingHealthCheck(ctx context.Context) ([]PoolKey, error) { + rows, err := q.db.QueryContext(ctx, listKeysNeedingHealthCheck) + if err != nil { + return nil, err + } + defer rows.Close() + items := []PoolKey{} + for rows.Next() { + var i PoolKey + if err := rows.Scan( + &i.ID, + &i.UserID, + &i.Label, + &i.KeyCiphertext, + &i.KeyFingerprint, + &i.Active, + &i.DailyLimitMicros, + &i.TodayDate, + &i.TodayMicros, + &i.TotalMicros, + &i.RequestCount, + &i.CreatedAt, + &i.LastUsedAt, + &i.MaxMicros, + &i.SharedMicros, + &i.PioneerTeamID, + &i.PioneerPaymentPlan, + &i.PioneerCreditLimitMicros, + &i.PioneerRemainingMicros, + &i.PioneerHealth, + &i.PioneerUnhealthySince, + &i.PendingValidation, + &i.LastBillingSyncAt, + &i.RevokedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const listUnhealthyKeysOlderThan = `-- name: ListUnhealthyKeysOlderThan :many +SELECT id, user_id, label, key_ciphertext, key_fingerprint, active, daily_limit_micros, today_date, today_micros, total_micros, request_count, created_at, last_used_at, max_micros, shared_micros, pioneer_team_id, pioneer_payment_plan, pioneer_credit_limit_micros, pioneer_remaining_micros, pioneer_health, pioneer_unhealthy_since, pending_validation, last_billing_sync_at, revoked_at FROM pool_keys +WHERE pioneer_health = 2 + AND pioneer_unhealthy_since IS NOT NULL + AND pioneer_unhealthy_since < ? + AND revoked_at IS NULL +` + +// Keys that have been unauthorized (pioneer_health=2) since before cutoff. +// Used by the reconciler to trigger permanent revocation after 14 days. +func (q *Queries) ListUnhealthyKeysOlderThan(ctx context.Context, pioneerUnhealthySince sql.NullInt64) ([]PoolKey, error) { + rows, err := q.db.QueryContext(ctx, listUnhealthyKeysOlderThan, pioneerUnhealthySince) + if err != nil { + return nil, err + } + defer rows.Close() + items := []PoolKey{} + for rows.Next() { + var i PoolKey + if err := rows.Scan( + &i.ID, + &i.UserID, + &i.Label, + &i.KeyCiphertext, + &i.KeyFingerprint, + &i.Active, + &i.DailyLimitMicros, + &i.TodayDate, + &i.TodayMicros, + &i.TotalMicros, + &i.RequestCount, + &i.CreatedAt, + &i.LastUsedAt, + &i.MaxMicros, + &i.SharedMicros, + &i.PioneerTeamID, + &i.PioneerPaymentPlan, + &i.PioneerCreditLimitMicros, + &i.PioneerRemainingMicros, + &i.PioneerHealth, + &i.PioneerUnhealthySince, + &i.PendingValidation, + &i.LastBillingSyncAt, + &i.RevokedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const markPoolKeyRevoked = `-- name: MarkPoolKeyRevoked :exec +UPDATE pool_keys SET + active = 0, + revoked_at = ? +WHERE id = ? +` + +type MarkPoolKeyRevokedParams struct { + RevokedAt sql.NullInt64 `json:"revoked_at"` + ID string `json:"id"` +} + +func (q *Queries) MarkPoolKeyRevoked(ctx context.Context, arg MarkPoolKeyRevokedParams) error { + _, err := q.db.ExecContext(ctx, markPoolKeyRevoked, arg.RevokedAt, arg.ID) + return err +} + +const pickPoolKeyV2 = `-- name: PickPoolKeyV2 :one +SELECT id, user_id, label, key_ciphertext, key_fingerprint, active, daily_limit_micros, today_date, today_micros, total_micros, request_count, created_at, last_used_at, max_micros, shared_micros, pioneer_team_id, pioneer_payment_plan, pioneer_credit_limit_micros, pioneer_remaining_micros, pioneer_health, pioneer_unhealthy_since, pending_validation, last_billing_sync_at, revoked_at FROM pool_keys +WHERE active = 1 + AND revoked_at IS NULL + AND pending_validation = 0 + AND pioneer_health = 1 + AND (pioneer_remaining_micros IS NULL OR pioneer_remaining_micros > 10000000) + AND today_micros < max_micros +ORDER BY today_micros ASC, RANDOM() +LIMIT 1 +` + +// Best active healthy key for a request: +// +// active=1, not revoked, not pending validation +// pioneer_health=1 (healthy) +// pioneer_remaining_micros > 10,000,000 ($10 buffer = 1000 credits) +// today_micros < max_micros +// +// Lowest today_micros wins; random tiebreak. +func (q *Queries) PickPoolKeyV2(ctx context.Context) (PoolKey, error) { + row := q.db.QueryRowContext(ctx, pickPoolKeyV2) + var i PoolKey + err := row.Scan( + &i.ID, + &i.UserID, + &i.Label, + &i.KeyCiphertext, + &i.KeyFingerprint, + &i.Active, + &i.DailyLimitMicros, + &i.TodayDate, + &i.TodayMicros, + &i.TotalMicros, + &i.RequestCount, + &i.CreatedAt, + &i.LastUsedAt, + &i.MaxMicros, + &i.SharedMicros, + &i.PioneerTeamID, + &i.PioneerPaymentPlan, + &i.PioneerCreditLimitMicros, + &i.PioneerRemainingMicros, + &i.PioneerHealth, + &i.PioneerUnhealthySince, + &i.PendingValidation, + &i.LastBillingSyncAt, + &i.RevokedAt, + ) + return i, err +} + +const updatePoolKeyHealth = `-- name: UpdatePoolKeyHealth :exec + +UPDATE pool_keys SET + pioneer_health = ?, + pioneer_unhealthy_since = ?, + pioneer_team_id = COALESCE(?, pioneer_team_id), + pioneer_payment_plan = COALESCE(?, pioneer_payment_plan), + pioneer_credit_limit_micros = COALESCE(?, pioneer_credit_limit_micros), + pioneer_remaining_micros = COALESCE(?, pioneer_remaining_micros), + today_micros = COALESCE(?, today_micros), + last_billing_sync_at = ? +WHERE id = ? +` + +type UpdatePoolKeyHealthParams struct { + PioneerHealth int64 `json:"pioneer_health"` + PioneerUnhealthySince sql.NullInt64 `json:"pioneer_unhealthy_since"` + PioneerTeamID sql.NullString `json:"pioneer_team_id"` + PioneerPaymentPlan sql.NullString `json:"pioneer_payment_plan"` + PioneerCreditLimitMicros sql.NullInt64 `json:"pioneer_credit_limit_micros"` + PioneerRemainingMicros sql.NullInt64 `json:"pioneer_remaining_micros"` + TodayMicros int64 `json:"today_micros"` + LastBillingSyncAt sql.NullInt64 `json:"last_billing_sync_at"` + ID string `json:"id"` +} + +// Pool key v2 queries: health tracking, billing sync, two-budget updates. +// +// pioneer_health integer enum: +// +// 0 = unknown +// 1 = healthy +// 2 = unauthorized +// +// Sets health and snapshots billing info from /plan-info. +// Pass NULL for optional fields to leave them unchanged. +func (q *Queries) UpdatePoolKeyHealth(ctx context.Context, arg UpdatePoolKeyHealthParams) error { + _, err := q.db.ExecContext(ctx, updatePoolKeyHealth, + arg.PioneerHealth, + arg.PioneerUnhealthySince, + arg.PioneerTeamID, + arg.PioneerPaymentPlan, + arg.PioneerCreditLimitMicros, + arg.PioneerRemainingMicros, + arg.TodayMicros, + arg.LastBillingSyncAt, + arg.ID, + ) + return err +} + +const updatePoolKeyLimits = `-- name: UpdatePoolKeyLimits :exec +UPDATE pool_keys SET + max_micros = ?, + shared_micros = ? +WHERE id = ? AND user_id = ? +` + +type UpdatePoolKeyLimitsParams struct { + MaxMicros int64 `json:"max_micros"` + SharedMicros int64 `json:"shared_micros"` + ID string `json:"id"` + UserID string `json:"user_id"` +} + +// Updates the two-budget limits. Server enforces 0 <= shared <= max. +func (q *Queries) UpdatePoolKeyLimits(ctx context.Context, arg UpdatePoolKeyLimitsParams) error { + _, err := q.db.ExecContext(ctx, updatePoolKeyLimits, + arg.MaxMicros, + arg.SharedMicros, + arg.ID, + arg.UserID, + ) + return err +} diff --git a/server/internal/store/potluck_requests.sql.go b/server/internal/store/potluck_requests.sql.go new file mode 100644 index 0000000..02aab7b --- /dev/null +++ b/server/internal/store/potluck_requests.sql.go @@ -0,0 +1,165 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 +// source: potluck_requests.sql + +package store + +import ( + "context" + "database/sql" +) + +const cancelPotluckRequest = `-- name: CancelPotluckRequest :exec +UPDATE potluck_requests SET + finished_at = ?, + status = 'canceled' +WHERE id = ? AND status = 'pending' +` + +type CancelPotluckRequestParams struct { + FinishedAt sql.NullInt64 `json:"finished_at"` + ID string `json:"id"` +} + +func (q *Queries) CancelPotluckRequest(ctx context.Context, arg CancelPotluckRequestParams) error { + _, err := q.db.ExecContext(ctx, cancelPotluckRequest, arg.FinishedAt, arg.ID) + return err +} + +const createPotluckRequest = `-- name: CreatePotluckRequest :one + +INSERT INTO potluck_requests ( + id, user_id, api_key_id, pool_key_id, surface, model, started_at, status +) VALUES (?, ?, ?, ?, ?, ?, ?, 'pending') +RETURNING id, user_id, api_key_id, pool_key_id, surface, model, started_at, finished_at, prompt_tokens, completion_tokens, total_tokens, status, error_code +` + +type CreatePotluckRequestParams struct { + ID string `json:"id"` + UserID string `json:"user_id"` + ApiKeyID sql.NullString `json:"api_key_id"` + PoolKeyID sql.NullString `json:"pool_key_id"` + Surface string `json:"surface"` + Model string `json:"model"` + StartedAt int64 `json:"started_at"` +} + +// potluck_requests: per-request log for attribution and billing join. +func (q *Queries) CreatePotluckRequest(ctx context.Context, arg CreatePotluckRequestParams) (PotluckRequest, error) { + row := q.db.QueryRowContext(ctx, createPotluckRequest, + arg.ID, + arg.UserID, + arg.ApiKeyID, + arg.PoolKeyID, + arg.Surface, + arg.Model, + arg.StartedAt, + ) + var i PotluckRequest + err := row.Scan( + &i.ID, + &i.UserID, + &i.ApiKeyID, + &i.PoolKeyID, + &i.Surface, + &i.Model, + &i.StartedAt, + &i.FinishedAt, + &i.PromptTokens, + &i.CompletionTokens, + &i.TotalTokens, + &i.Status, + &i.ErrorCode, + ) + return i, err +} + +const finishPotluckRequest = `-- name: FinishPotluckRequest :exec +UPDATE potluck_requests SET + finished_at = ?, + prompt_tokens = ?, + completion_tokens = ?, + total_tokens = ?, + status = ? +WHERE id = ? +` + +type FinishPotluckRequestParams struct { + FinishedAt sql.NullInt64 `json:"finished_at"` + PromptTokens sql.NullInt64 `json:"prompt_tokens"` + CompletionTokens sql.NullInt64 `json:"completion_tokens"` + TotalTokens sql.NullInt64 `json:"total_tokens"` + Status string `json:"status"` + ID string `json:"id"` +} + +func (q *Queries) FinishPotluckRequest(ctx context.Context, arg FinishPotluckRequestParams) error { + _, err := q.db.ExecContext(ctx, finishPotluckRequest, + arg.FinishedAt, + arg.PromptTokens, + arg.CompletionTokens, + arg.TotalTokens, + arg.Status, + arg.ID, + ) + return err +} + +const listUnmatchedRequestsForKey = `-- name: ListUnmatchedRequestsForKey :many +SELECT pr.id, pr.user_id, pr.api_key_id, pr.pool_key_id, pr.surface, pr.model, pr.started_at, pr.finished_at, pr.prompt_tokens, pr.completion_tokens, pr.total_tokens, pr.status, pr.error_code FROM potluck_requests pr +WHERE pr.pool_key_id = ? + AND pr.finished_at >= ? + AND pr.finished_at <= ? + AND pr.status = 'done' + AND pr.id NOT IN ( + SELECT pkbr.matched_request_id FROM pool_key_billing_rows pkbr + WHERE pkbr.matched_request_id IS NOT NULL + ) +ORDER BY pr.finished_at ASC +` + +type ListUnmatchedRequestsForKeyParams struct { + PoolKeyID sql.NullString `json:"pool_key_id"` + FinishedAt sql.NullInt64 `json:"finished_at"` + FinishedAt_2 sql.NullInt64 `json:"finished_at_2"` +} + +// Requests that finished in a given time window with no billing row matched yet. +// Used by the reconciler attribution pass. +func (q *Queries) ListUnmatchedRequestsForKey(ctx context.Context, arg ListUnmatchedRequestsForKeyParams) ([]PotluckRequest, error) { + rows, err := q.db.QueryContext(ctx, listUnmatchedRequestsForKey, arg.PoolKeyID, arg.FinishedAt, arg.FinishedAt_2) + if err != nil { + return nil, err + } + defer rows.Close() + items := []PotluckRequest{} + for rows.Next() { + var i PotluckRequest + if err := rows.Scan( + &i.ID, + &i.UserID, + &i.ApiKeyID, + &i.PoolKeyID, + &i.Surface, + &i.Model, + &i.StartedAt, + &i.FinishedAt, + &i.PromptTokens, + &i.CompletionTokens, + &i.TotalTokens, + &i.Status, + &i.ErrorCode, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} diff --git a/server/internal/store/querier.go b/server/internal/store/querier.go index 4320633..1c6d470 100644 --- a/server/internal/store/querier.go +++ b/server/internal/store/querier.go @@ -10,10 +10,14 @@ import ( ) type Querier interface { + // Called by the reconciler when a pending_validation or unauthorized key + // comes back healthy (pioneer_health=1). + ActivatePoolKey(ctx context.Context, id string) error AddAPIKeySpend(ctx context.Context, arg AddAPIKeySpendParams) error AppendAssistantContent(ctx context.Context, arg AppendAssistantContentParams) error AppendStreamChunk(ctx context.Context, arg AppendStreamChunkParams) error ArchiveConversation(ctx context.Context, arg ArchiveConversationParams) error + CancelPotluckRequest(ctx context.Context, arg CancelPotluckRequestParams) error CountActiveStreamsForUser(ctx context.Context, userID string) (int64, error) CreateAPIKey(ctx context.Context, arg CreateAPIKeyParams) (ApiKey, error) CreateContribution(ctx context.Context, arg CreateContributionParams) (Contribution, error) @@ -25,6 +29,8 @@ type Querier interface { // under its daily cap. Spend counters are maintained here; exact accounting // lives in the spends table. CreatePoolKey(ctx context.Context, arg CreatePoolKeyParams) (PoolKey, error) + // potluck_requests: per-request log for attribution and billing join. + CreatePotluckRequest(ctx context.Context, arg CreatePotluckRequestParams) (PotluckRequest, error) CreateSession(ctx context.Context, arg CreateSessionParams) (Session, error) CreateStream(ctx context.Context, arg CreateStreamParams) (Stream, error) CreateUser(ctx context.Context, arg CreateUserParams) (User, error) @@ -33,10 +39,15 @@ type Querier interface { DeletePoolKey(ctx context.Context, arg DeletePoolKeyParams) error DeleteSession(ctx context.Context, id string) error DeleteSessionForUser(ctx context.Context, arg DeleteSessionForUserParams) error + FinishPotluckRequest(ctx context.Context, arg FinishPotluckRequestParams) error GetAPIKeyByHash(ctx context.Context, keyHash string) (ApiKey, error) GetConversation(ctx context.Context, arg GetConversationParams) (Conversation, error) GetIdempotency(ctx context.Context, arg GetIdempotencyParams) (IdempotencyKey, error) + // Most recent recompute for today, for display in the dashboard. + GetLatestRecompute(ctx context.Context, day int64) (GetLatestRecomputeRow, error) GetMessage(ctx context.Context, id string) (Message, error) + // Oldest refreshed_at across all models — tells us when the catalog is stale. + GetModelCatalogRefreshedAt(ctx context.Context) (interface{}, error) GetModelPrice(ctx context.Context, model string) (ModelPrice, error) GetPoolKey(ctx context.Context, id string) (PoolKey, error) GetSession(ctx context.Context, arg GetSessionParams) (Session, error) @@ -45,11 +56,26 @@ type Querier interface { GetUserByEmail(ctx context.Context, email string) (User, error) GetUserByHCAID(ctx context.Context, hcaID sql.NullString) (User, error) GetUserByID(ctx context.Context, id string) (User, error) + GetUserDailyAllowance(ctx context.Context, arg GetUserDailyAllowanceParams) (UserDailyAllowance, error) + GetUserDailySpend(ctx context.Context, arg GetUserDailySpendParams) (UserDailySpend, error) + // Most recent pioneer_created_at for a key. + // COALESCE(SUM(pioneer_created_at*0) + MAX(...)) is a workaround for + // sqlc's sqlite parser rejecting bare COALESCE(MAX(...), 0). + LatestBillingRowTime(ctx context.Context, poolKeyID string) (interface{}, error) ListAPIKeysForUser(ctx context.Context, userID string) ([]ApiKey, error) + // All billing rows for a key after a given timestamp. + ListBillingRowsForKeyAfter(ctx context.Context, arg ListBillingRowsForKeyAfterParams) ([]PoolKeyBillingRow, error) + // Billing rows attributed to a user in a time window. + // Caller filters duplicates and sums in Go. + ListBillingRowsForUserToday(ctx context.Context, arg ListBillingRowsForUserTodayParams) ([]PoolKeyBillingRow, error) ListContributionsForUser(ctx context.Context, arg ListContributionsForUserParams) ([]Contribution, error) ListConversationsForUser(ctx context.Context, arg ListConversationsForUserParams) ([]Conversation, error) ListEstimatedSpends(ctx context.Context, limit int64) ([]Spend, error) + // Keys the reconciler should probe on each tick. + // Excludes permanently revoked keys. + ListKeysNeedingHealthCheck(ctx context.Context) ([]PoolKey, error) ListMessagesForConversation(ctx context.Context, conversationID string) ([]Message, error) + ListModelCatalog(ctx context.Context) ([]ModelsCatalog, error) ListModelPrices(ctx context.Context) ([]ModelPrice, error) // Per-model aggregate: all-time token/spend totals + 48h TPS average. // The since parameter scopes only the TPS calculation; token counts are @@ -63,15 +89,31 @@ type Querier interface { ListPoolKeysForUser(ctx context.Context, userID string) ([]PoolKey, error) ListSessionsForUser(ctx context.Context, arg ListSessionsForUserParams) ([]Session, error) ListStreamChunksAfter(ctx context.Context, arg ListStreamChunksAfterParams) ([]ListStreamChunksAfterRow, error) + // Keys that have been unauthorized (pioneer_health=2) since before cutoff. + // Used by the reconciler to trigger permanent revocation after 14 days. + ListUnhealthyKeysOlderThan(ctx context.Context, pioneerUnhealthySince sql.NullInt64) ([]PoolKey, error) + // Requests that finished in a given time window with no billing row matched yet. + // Used by the reconciler attribution pass. + ListUnmatchedRequestsForKey(ctx context.Context, arg ListUnmatchedRequestsForKeyParams) ([]PotluckRequest, error) // Per-user contribution totals, spend totals, and derived balance. // Used by the allocation calculator on the dashboard. ListUserAllocations(ctx context.Context) ([]ListUserAllocationsRow, error) + ListUserDailyAllowancesForDay(ctx context.Context, day int64) ([]UserDailyAllowance, error) + ListUserDailySpendForDay(ctx context.Context, day int64) ([]UserDailySpend, error) + MarkPoolKeyRevoked(ctx context.Context, arg MarkPoolKeyRevokedParams) error MaxStreamChunkSeq(ctx context.Context, streamID string) (interface{}, error) // Select the best key to use for a request: active, under daily cap, // least spend today. Resets stale today_* counters are handled in Go // (compare today_date to current UTC day and update if stale). // ?1 = current UTC day (unix / 86400) PickPoolKey(ctx context.Context, todayDate int64) (PoolKey, error) + // Best active healthy key for a request: + // active=1, not revoked, not pending validation + // pioneer_health=1 (healthy) + // pioneer_remaining_micros > 10,000,000 ($10 buffer = 1000 credits) + // today_micros < max_micros + // Lowest today_micros wins; random tiebreak. + PickPoolKeyV2(ctx context.Context) (PoolKey, error) PoolActiveKeyCount(ctx context.Context) (int64, error) PoolContributorCount(ctx context.Context) (int64, error) PoolSpentSince(ctx context.Context, createdAt int64) (interface{}, error) @@ -104,15 +146,39 @@ type Querier interface { TouchSession(ctx context.Context, arg TouchSessionParams) error TouchUser(ctx context.Context, arg TouchUserParams) error UpdateConversationTitle(ctx context.Context, arg UpdateConversationTitleParams) error + // Pool key v2 queries: health tracking, billing sync, two-budget updates. + // + // pioneer_health integer enum: + // 0 = unknown + // 1 = healthy + // 2 = unauthorized + // Sets health and snapshots billing info from /plan-info. + // Pass NULL for optional fields to leave them unchanged. + UpdatePoolKeyHealth(ctx context.Context, arg UpdatePoolKeyHealthParams) error UpdatePoolKeyLabel(ctx context.Context, arg UpdatePoolKeyLabelParams) error UpdatePoolKeyLimit(ctx context.Context, arg UpdatePoolKeyLimitParams) error + // Updates the two-budget limits. Server enforces 0 <= shared <= max. + UpdatePoolKeyLimits(ctx context.Context, arg UpdatePoolKeyLimitsParams) error + // pool_key_billing_rows: ingested pioneer billing log entries. + // + // attribution integer enum: + // 0 = matched (joined to a potluck_requests row) + // 1 = judge_paired (/llmaj/judge paired to preceding opus call) + // 2 = owner_fallback (no match, charged to key owner) + // 3 = duplicate (double-logged by pioneer, not charged) + UpsertBillingRow(ctx context.Context, arg UpsertBillingRowParams) error UpsertMessage(ctx context.Context, arg UpsertMessageParams) (Message, error) + // models_catalog: hourly-refreshed cache of /v1/models + /base-models. + UpsertModelCatalog(ctx context.Context, arg UpsertModelCatalogParams) error UpsertModelPrice(ctx context.Context, arg UpsertModelPriceParams) error UpsertSpend(ctx context.Context, arg UpsertSpendParams) (Spend, error) // Find-or-create by HCA id, refreshing the cached identity fields on each // successful sign-in. Email is updated too because HCA users can change // theirs and the local copy should track upstream. UpsertUserByHCAID(ctx context.Context, arg UpsertUserByHCAIDParams) (User, error) + UpsertUserDailyAllowance(ctx context.Context, arg UpsertUserDailyAllowanceParams) error + // user_daily_spend and user_daily_allowances. + UpsertUserDailySpend(ctx context.Context, arg UpsertUserDailySpendParams) error } var _ Querier = (*Queries)(nil)