diff options
| author | Chia <Chia@93.nz> | 2026-08-05 00:26:25 +1200 |
|---|---|---|
| committer | Chia <Chia@93.nz> | 2026-08-05 00:33:31 +1200 |
| commit | 1a3d7f9a8a181df48f0e911cbe17a3fad3ab9ac9 (patch) | |
| tree | 8c92e1e7326fc67ed077a0a878697f1be14b43da | |
| parent | 5b651488b081b65fda8a323f228e139adb79a35d (diff) | |
add some scriptsmain
Diffstat (limited to '')
41 files changed, 3719 insertions, 161 deletions
diff --git a/.env.control.example b/.env.control.example index 4654d0c..8152d40 100644 --- a/.env.control.example +++ b/.env.control.example @@ -3,4 +3,6 @@ AIGW_REDIS_URL=redis://127.0.0.1:6379/0 # base64 of exactly 32 random bytes; generate with: openssl rand -base64 32 AIGW_CREDENTIAL_KEY=replace-with-base64-32-byte-key AIGW_ADMIN_TOKEN=replace-with-a-long-random-admin-token - +# Prefer a restricted test key (rk_test_) with Checkout Session write access. +AIGW_STRIPE_API_KEY=replace-with-a-stripe-restricted-key +AIGW_STRIPE_WEBHOOK_SECRET=replace-with-a-webhook-signing-secret @@ -1,5 +1,6 @@ .env -aigw +.env.debug +/aigw coverage.out .playwright-cli/ output/ @@ -15,6 +15,11 @@ AIGW 是一个轻量、无状态的 AI API 中转后端。当前阶段聚焦上 - `Authorization: Bearer` 和 `x-api-key` 客户鉴权 - 统一 JSON 错误、`X-AIGW-Request-ID`、Prometheus 文本指标 - 非阻塞用量事件，包含租户、项目、模型、上游、尝试次数、耗时和 token 用量 +- PostgreSQL 预付余额、请求额度冻结、实际 token 结算和不可变账本 +- Stripe 托管 Checkout 充值、签名 Webhook 与事件/订单双重幂等 +- PostgreSQL Usage Ledger、月度项目汇总和单请求成本追溯 +- 平台/租户控制台令牌、六种 RBAC 角色和管理 API 审计日志 +- 项目级 RPM、估算 TPM、并发限制和月度消费配额 ## 快速运行 @@ -83,12 +88,13 @@ curl http://127.0.0.1:8080/anthropic/v1/messages \ ## 生产边界 -当前版本可以作为第一阶段的数据面，但还不是完整商业计费平台： +当前版本可以作为带预付计费的数据面，但上线前仍需完成业务侧对账： -- 用量事件写到结构化日志，队列满时会丢弃并增加 `aigw_usage_events_dropped_total`。正式扣费前必须替换为持久、幂等的账本写入或消息队列。 +- 控制面计费模式会把 UsageEvent、冻结记录和扣费流水同步、幂等写入 PostgreSQL；异步结构化日志只是可观测副本，丢弃不会影响账本。 - 客户密钥和模型目录在控制面模式下从 PostgreSQL 载入到原子内存快照；Redis 只是可选的变更广播加速层，故障时通过 PostgreSQL generation 轮询收敛。 - 当前只把 OpenAI 入口发给 OpenAI 兼容上游、Anthropic 入口发给 Anthropic 兼容上游，不做跨协议转换。 - 自动故障转移可能在极少数网络错误下造成上游重复执行。正式计费时需要上游幂等能力、请求去重策略和重复成本对账。 +- 上游必须返回 usage 才能按 token 结算；未返回 usage 的成功响应当前记为零费用，应在上生产前为每个供应商做账单对账或补充 token 计算器。 - `/metrics` 应仅在内网暴露；公网 TLS、WAF 和连接层限速应放在负载均衡器或边缘代理。 ## PostgreSQL + Redis 控制面 @@ -100,12 +106,60 @@ curl http://127.0.0.1:8080/anthropic/v1/messages \ 开发环境可以直接启动： ```bash -export AIGW_CREDENTIAL_KEY="$(openssl rand -base64 32)" -export AIGW_ADMIN_TOKEN="$(openssl rand -hex 32)" +./scripts/start-debug.sh +``` + +脚本会首次生成仅当前用户可读的 `.env.debug`，构建并等待 PostgreSQL、Redis 和网关健康，然后打印后台地址和管理员 token。关闭调试服务： + +```bash +./scripts/stop-debug.sh +``` + +关闭脚本保留 PostgreSQL/Redis 数据卷，下一次启动仍可继续使用已有控制面数据。需要测试 Stripe Checkout 时，先把 `.env.debug` 中两个 Stripe 占位值替换成测试环境 restricted key 和 Webhook signing secret。 + +源码未变化时可跳过镜像构建以快速重启：`AIGW_DEBUG_SKIP_BUILD=1 ./scripts/start-debug.sh`。默认构建使用 Docker host network；特殊环境可以通过 `AIGW_DOCKER_BUILD_NETWORK=default` 覆盖。 + +然后打开 `http://127.0.0.1:8080/admin/`，输入脚本打印的 `AIGW_ADMIN_TOKEN`。这个 token 是平台管理员的 bootstrap/break-glass 凭证；日常操作应在 Team 页面签发数据库控制台令牌。第一套资源的创建顺序是：Tenant → Project → API key → Provider → Model route。客户 API Key 和控制台令牌的明文都只在创建成功时返回一次。 + +控制台角色分为：`platform_admin`、`platform_viewer`、`tenant_admin`、`tenant_billing`、`tenant_developer`、`tenant_viewer`。租户角色的查询条件在服务端下推到 PostgreSQL，不能读取其他租户的项目、密钥、余额、Usage 或审计事件；供应商凭证和路由管理只对平台角色开放。 + +## Usage、配额与限流 + +每个完成上游尝试的请求都会按 `request_id` 幂等写入 `usage_events`，并更新 `usage_monthly_rollups`。Usage 持久化独立于预付费冻结记录，因此关闭计费也不会关闭用量账本。Admin WebUI 提供本月汇总、单请求状态、模型、token、成本、未收金额和延迟查询。 + +Limits 页面按项目配置： + +- `requests_per_minute`：固定分钟窗口请求数，`0` 表示不限。 +- `tokens_per_minute`：请求体约算输入 token 加显式/默认最大输出 token，`0` 表示不限。 +- `concurrent_requests`：从进入上游前到响应复制完成的并发 lease，`0` 表示不限。 +- `monthly_spend_micros`：PG 事务内检查当月已计成本和未结冻结，`0` 表示不限。 + +Redis 可用时，RPM/TPM/并发通过 Lua 原子执行并在多实例间共享；Redis 故障时自动退回本机计数，服务继续可用，但降级期间限制是“每实例”而不是“全局”。Redis 恢复后新请求会自动重新使用分布式计数。启用预付计费时，月度消费配额和余额由 PostgreSQL 保证，不依赖 Redis。 + +## 余额与 Stripe 充值 + +模型价格在后台按“币种单位 / 100 万 token”配置，数据库使用 `amount_micros` 固定精度整数保存金额。推理请求会先按请求体字节数和 `max_tokens`/`max_completion_tokens` 保守冻结额度；成功响应按上游返回的输入、输出和缓存 token 结算，失败请求释放冻结。`request_id` 是用量与扣费幂等键。 + +Stripe 使用托管 Checkout，服务端不会接触卡号，也没有硬编码支付方式；支付方式由 Stripe Dashboard 动态配置。充值只在签名校验通过的 Webhook 确认 `payment_status=paid` 后入账，成功跳转页不会直接修改余额。 + +配置 Stripe 测试环境： + +```bash +cp .env.control.example .env +# 将 AIGW_STRIPE_API_KEY 设为最小权限的 rk_test_ restricted key +# 本地转发会显示 whsec_...，填入 AIGW_STRIPE_WEBHOOK_SECRET +stripe listen --forward-to http://127.0.0.1:8080/billing/stripe/webhook docker compose up --build ``` -然后打开 `http://127.0.0.1:8080/admin/`，输入 `AIGW_ADMIN_TOKEN`。第一套资源的创建顺序是：Tenant → Project → API key → Provider → Model route。创建客户 API Key 时，明文只在成功响应中出现一次。 +Webhook 至少订阅： + +- `checkout.session.completed` +- `checkout.session.async_payment_succeeded` +- `checkout.session.async_payment_failed` +- `checkout.session.expired` + +当前没有默认启用 Stripe Tax，因为是否有有效税务注册不能由代码推断。确认注册和税务处理方案后再显式加入 `automatic_tax`。生产环境应把 Stripe restricted key 和 Webhook signing secret 放入云平台的密钥管理服务，并限制密钥权限和来源 IP，不要放进镜像或仓库。 生产环境建议把 `auto_migrate` 改为 `false`，先执行： @@ -113,7 +167,7 @@ docker compose up --build AIGW_DATABASE_URL="postgres://..." go run ./cmd/migrate ``` -管理 API 仅使用一个 bootstrap token，应该放在内网、VPN 或反向代理后。下一阶段需要把它替换为管理员用户、角色和审计日志。 +管理 API 支持 bootstrap token 和数据库控制台令牌。即使已经启用 RBAC，管理监听端口仍应放在内网、VPN 或身份感知反向代理后；bootstrap token 应只用于首次建号和故障恢复。 完整的扩展边界见 [架构说明](docs/architecture.md)。 diff --git a/cmd/aigw/main.go b/cmd/aigw/main.go new file mode 100644 index 0000000..de929de --- /dev/null +++ b/cmd/aigw/main.go @@ -0,0 +1,186 @@ +package main + +import ( + "context" + "errors" + "flag" + "fmt" + "log/slog" + "net/http" + "os" + "os/signal" + "syscall" + "time" + + "aigw/internal/adminapi" + "aigw/internal/auth" + "aigw/internal/billing" + "aigw/internal/catalog" + "aigw/internal/config" + "aigw/internal/controlplane" + "aigw/internal/httpapi" + "aigw/internal/limits" + "aigw/internal/provider" + "aigw/internal/routing" + "aigw/internal/telemetry" +) + +func main() { + configPath := flag.String("config", "config.json", "path to the gateway configuration") + flag.Parse() + logger := slog.New(slog.NewJSONHandler(os.Stdout, nil)) + cfg, err := config.Load(*configPath) + if err != nil { + logger.Error("configuration_failed", "error", err) + os.Exit(1) + } + ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + if err := run(ctx, cfg, logger); err != nil { + logger.Error("gateway_stopped", "error", err) + os.Exit(1) + } +} + +func run(ctx context.Context, cfg config.Config, logger *slog.Logger) error { + metrics := &telemetry.Metrics{} + usageSink := telemetry.NewAsyncUsageLogger(logger, metrics, cfg.Observability.UsageBuffer) + defer func() { + closeContext, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := usageSink.Close(closeContext); err != nil { + logger.Warn("usage_sink_close_failed", "error", err) + } + }() + + var authenticator *auth.StaticAuthenticator + var modelCatalog *catalog.Catalog + var store *controlplane.Store + var manager *controlplane.Manager + var managerCancel context.CancelFunc + var managerDone chan struct{} + var requestLimiter *limits.Limiter + if cfg.ControlPlane.Enabled { + var err error + store, err = controlplane.NewStore(ctx, controlplane.Options{ + DatabaseURL: cfg.ControlPlane.DatabaseURL, RedisURL: cfg.ControlPlane.RedisURL, + CredentialKey: cfg.ControlPlane.CredentialKey, RedisChannel: cfg.ControlPlane.RedisChannel, + VersionCacheKey: cfg.ControlPlane.SnapshotCacheKey, + }) + if err != nil { + return err + } + defer store.Close() + if cfg.ControlPlane.AutoMigrate { + if err := store.Migrate(ctx); err != nil { + return err + } + } + authenticator = auth.NewDynamic(nil, cfg.Auth.AllowAnonymous) + modelCatalog = catalog.NewModels(nil) + requestLimiter = limits.New(cfg.ControlPlane.RedisURL, "aigw:limits", cfg.Billing.DefaultMaxOutputTokens, logger) + defer requestLimiter.Close() + manager = controlplane.NewManager(store, modelCatalog, authenticator, logger, + time.Duration(cfg.ControlPlane.ReloadIntervalSeconds)*time.Second, requestLimiter) + if _, err := manager.Reload(ctx); err != nil { + return fmt.Errorf("load initial control-plane snapshot: %w", err) + } + managerContext, cancel := context.WithCancel(ctx) + managerCancel = cancel + managerDone = make(chan struct{}) + go func() { + defer close(managerDone) + manager.Run(managerContext) + }() + defer func() { + managerCancel() + select { + case <-managerDone: + case <-time.After(3 * time.Second): + logger.Warn("control_plane_shutdown_timed_out") + } + }() + } else { + var err error + authenticator, err = auth.NewStatic(os.Getenv(cfg.Auth.KeysEnv), cfg.Auth.AllowAnonymous) + if err != nil { + return err + } + modelCatalog = catalog.New(cfg) + } + + var billingService *billing.Service + if cfg.Billing.Enabled { + var err error + billingService, err = billing.New(ctx, billing.Options{ + DatabaseURL: cfg.ControlPlane.DatabaseURL, Currency: cfg.Billing.Currency, + DefaultMaxOutputTokens: cfg.Billing.DefaultMaxOutputTokens, + MinTopUpMinor: cfg.Billing.MinTopUpMinor, MaxTopUpMinor: cfg.Billing.MaxTopUpMinor, + StripeEnabled: cfg.Billing.Stripe.Enabled, StripeAPIKey: cfg.Billing.Stripe.APIKey, + StripeWebhookSecret: cfg.Billing.Stripe.WebhookSecret, + StripeSuccessURL: cfg.Billing.Stripe.SuccessURL, StripeCancelURL: cfg.Billing.Stripe.CancelURL, + }) + if err != nil { + return err + } + defer billingService.Close() + } + var billingMeter billing.Meter + if billingService != nil { + billingMeter = billingService + } + + inferenceAPI := httpapi.New(httpapi.Options{ + Authenticator: authenticator, + Catalog: modelCatalog, + Router: routing.New(modelCatalog), + Forwarder: provider.New(cfg.UpstreamHTTP, metrics), + UsageSink: usageSink, + BillingMeter: billingMeter, + Limiter: requestLimiter, + UsageRecorder: store, + Metrics: metrics, + Logger: logger, + MaxBodyBytes: cfg.Server.MaxBodyBytes, + ExposeMetrics: cfg.Observability.ExposeMetrics, + }) + root := http.NewServeMux() + root.Handle("/", inferenceAPI.Handler()) + if cfg.Admin.Enabled { + adminHandler := adminapi.New(adminapi.Options{ + Store: store, Manager: manager, Billing: billingService, Token: cfg.Admin.Token, + Logger: logger, Prefix: cfg.Admin.BasePath, + }).Handler() + root.Handle(cfg.Admin.BasePath, adminHandler) + root.Handle(cfg.Admin.BasePath+"/", adminHandler) + } + if billingService != nil && billingService.StripeEnabled() { + root.Handle("/billing/stripe/webhook", billingService.WebhookHandler()) + } + + server := &http.Server{ + Addr: cfg.Server.Address, Handler: root, + ReadHeaderTimeout: cfg.Server.ReadHeaderTimeout(), + IdleTimeout: cfg.Server.IdleTimeout(), + } + errCh := make(chan error, 1) + go func() { + logger.Info("gateway_listening", "address", cfg.Server.Address, "control_plane", cfg.ControlPlane.Enabled, + "billing", cfg.Billing.Enabled, "stripe", cfg.Billing.Stripe.Enabled) + errCh <- server.ListenAndServe() + }() + select { + case <-ctx.Done(): + shutdownContext, cancel := context.WithTimeout(context.Background(), cfg.Server.ShutdownTimeout()) + defer cancel() + if err := server.Shutdown(shutdownContext); err != nil { + return fmt.Errorf("shutdown HTTP server: %w", err) + } + return nil + case err := <-errCh: + if errors.Is(err, http.ErrServerClosed) { + return nil + } + return fmt.Errorf("serve HTTP: %w", err) + } +} diff --git a/config.control.example.json b/config.control.example.json index d2affaf..d601b43 100644 --- a/config.control.example.json +++ b/config.control.example.json @@ -24,6 +24,20 @@ "token_env": "AIGW_ADMIN_TOKEN", "base_path": "/admin" }, + "billing": { + "enabled": true, + "currency": "usd", + "default_max_output_tokens": 4096, + "min_top_up_minor": 500, + "max_top_up_minor": 1000000, + "stripe": { + "enabled": true, + "api_key_env": "AIGW_STRIPE_API_KEY", + "webhook_secret_env": "AIGW_STRIPE_WEBHOOK_SECRET", + "success_url": "http://127.0.0.1:8080/admin/?topup=success", + "cancel_url": "http://127.0.0.1:8080/admin/?topup=cancel" + } + }, "upstream_http": { "max_idle_connections": 4096, "max_idle_connections_per_host": 1024, @@ -37,4 +51,3 @@ "providers": [], "models": [] } - diff --git a/docker-compose.yml b/docker-compose.yml index 9e74612..e2dde2a 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -28,6 +28,7 @@ services: retries: 20 aigw: + image: aigw-debug:local build: . depends_on: postgres: @@ -39,8 +40,16 @@ services: AIGW_REDIS_URL: redis://redis:6379/0 AIGW_CREDENTIAL_KEY: ${AIGW_CREDENTIAL_KEY:?set AIGW_CREDENTIAL_KEY} AIGW_ADMIN_TOKEN: ${AIGW_ADMIN_TOKEN:?set AIGW_ADMIN_TOKEN} + AIGW_STRIPE_API_KEY: ${AIGW_STRIPE_API_KEY:?set AIGW_STRIPE_API_KEY} + AIGW_STRIPE_WEBHOOK_SECRET: ${AIGW_STRIPE_WEBHOOK_SECRET:?set AIGW_STRIPE_WEBHOOK_SECRET} volumes: - ./config.control.example.json:/etc/aigw/config.json:ro + healthcheck: + test: ["CMD", "wget", "-q", "-O", "-", "http://127.0.0.1:8080/readyz"] + interval: 3s + timeout: 2s + retries: 20 + start_period: 5s volumes: aigw-postgres: diff --git a/docs/architecture.md b/docs/architecture.md index 640a37d..e522dd5 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -2,20 +2,24 @@ ## 设计目标 -第一阶段只做高性能、可横向扩容的数据面，并提前固定未来计费、权限和控制面的接口位置。代理核心不在请求热路径同步访问数据库，也不把供应商差异扩散到对外 API 层。 +网关保持轻量且可横向扩容，并把计费、权限和控制面边界固定下来。路由和鉴权只读内存快照；启用预付计费时，请求热路径会执行余额冻结与结算事务，以财务正确性优先。 ```mermaid flowchart LR C["客户 SDK"] --> E["OpenAI / Anthropic 入口"] E --> A["Authenticator + Scope"] - A --> R["Model Catalog + Router"] + A --> Q["Redis/local limits"] + Q --> B["Balance + monthly quota"] + B --> R["Model Catalog + Router"] R --> P["共享连接池 + Provider Adapter"] P --> U1["上游 A"] P --> U2["上游 B"] P --> U3["后备上游"] - E -. "非阻塞" .-> M["Usage Event Sink"] + P --> M["Usage settlement"] E -.-> O["Metrics + Structured Logs"] - M --> L["未来：持久计费账本"] + M --> L["PostgreSQL usage + ledger"] + S1["Stripe Checkout"] --> W["Signed Webhook"] + W --> L D["PostgreSQL 控制面"] -. "事实数据" .-> S["Redis generation / PubSub"] S -. "变更广播" .-> A D -. "启动与轮询快照" .-> A @@ -25,29 +29,33 @@ flowchart LR ## 热路径 1. 入口生成不可预测的请求 ID，并通过 Bearer 或 `x-api-key` 解析客户身份。 -2. 身份包含 `key_id`、`tenant_id`、`project_id` 和 scopes；当前是静态实现，代理只依赖 `Authenticator` 接口。 +2. 身份包含 `key_id`、`tenant_id`、`project_id` 和 scopes；控制面模式使用 PostgreSQL 摘要快照，代理只依赖 `Authenticator` 接口。 3. 请求体在配置上限内读取一次，以提取公开模型名并支持故障转移时重放。 -4. Router 按协议过滤路由，先按优先级分组，再在同级内按权重选择首选上游。 -5. Provider Adapter 重写上游模型名和凭证，使用进程级共享 Transport 发送请求。 -6. 上游返回成功头后即锁定路由。SSE 使用固定缓冲区复制并逐块 flush，不等待完整响应。 -7. 请求结束后发布 UsageEvent；发布操作有界、非阻塞，当前消费者输出结构化日志。 +4. 项目策略从 PG 热更新快照读取。Redis Lua 原子占用 RPM、估算 TPM 和并发额度；Redis 不可用时退回本机窗口。 +5. Router 按协议过滤路由，先按优先级分组，再在同级内按权重选择首选上游。 +6. 计费开启时，请求进入上游前在 PostgreSQL 原子检查月度消费、冻结保守估算额度；余额不足返回 402，月度额度耗尽返回 429。 +7. Provider Adapter 重写上游模型名和凭证，使用进程级共享 Transport 发送请求。上游返回成功头后即锁定路由；SSE 逐块 flush。 +8. 请求结束后释放并发 lease，并用 `request_id` 幂等持久化 UsageEvent、更新月度汇总、按实际 token 结算；结构化日志仅是异步副本。 ## 已固定的扩展边界 | 边界 | 当前实现 | 下一阶段替换 | | --- | --- | --- | -| 客户身份 | PostgreSQL 快照、内存 SHA-256 索引 | 管理员 RBAC、撤销审计、Redis/Lua 配额 | -| 权限 | `Principal.Scopes` 中的 `inference` | RBAC/ABAC、模型 allowlist、IP 与预算策略 | +| 客户身份 | PostgreSQL 快照、内存 SHA-256 索引 | SSO/OIDC、SCIM、模型 allowlist | +| 控制台权限 | 平台/租户令牌、六角色 RBAC、租户 SQL scope、审计日志 | SSO、细粒度自定义角色、审批流 | +| 权限 | `Principal.Scopes` 中的 `inference` + 控制台 RBAC | ABAC、IP 与模型策略 | | 模型目录 | PostgreSQL 快照 + 可选 Redis generation 广播 + PG 轮询兜底 | 版本化控制面、热更新、灰度发布 | | 路由 | priority + weighted selection + failover | 健康评分、延迟 EWMA、成本/质量策略、熔断 | -| 计量 | UsageEvent 异步日志 | 持久消息流、幂等消费、不可变用量账本 | -| 计费 | 未执行扣费 | 价格版本、预授权/额度、结算、退款与对账 | -| 限流 | 仅上游 429 故障转移 | Redis/Lua 或边缘 token bucket、租户并发额度 | +| 计量 | PostgreSQL UsageEvent + 月度汇总 + 异步日志副本 | 分区表、持久消息流、供应商账单对账 | +| 计费 | 预付余额、冻结/结算、不可变流水、Stripe Checkout | 价格版本、退款/争议、信用额度、Metronome 企业合同 | +| 限流 | Redis Lua 全局 RPM/估算 TPM/并发，故障时本机降级；PG 月度消费配额 | 滑动窗口、层级策略、边缘 token bucket | | 协议 | 同协议透传 | 规范化 IR + OpenAI/Anthropic/Google 双向转换 | ## 计费数据原则 -真实计费不能直接依赖请求日志。建议下一阶段建立不可变 `usage_ledger`，以 `request_id + attempt` 做幂等键，并同时记录：公开模型、实际上游、价格表版本、输入/输出/cache token、币种、上游成本、客户价格、状态和冲正关系。 +真实计费不依赖请求日志。当前以 `request_id` 作为 reservation、usage 和扣费幂等键，价格快照随冻结记录保存，账本流水不可变。余额使用百万分之一币种单位，所有变更在锁定 tenant wallet 的 PostgreSQL 事务中完成。 + +Stripe 充值使用 Checkout Session：本地先创建 top-up order，Stripe 请求使用 order ID 作为幂等键；Webhook 验证签名后再次核对 event ID、order ID、session ID、金额和币种。浏览器成功跳转不具有入账权威性。 流式请求的 usage 可能只在最后事件出现。当前 observer 会在线解析 OpenAI/Anthropic SSE 中的 usage；如果上游不返回 usage，事件中的 token 为零。接入商业扣费前，应为每种上游建立有测试的 usage normalizer，并使用上游账单进行日对账。 @@ -57,12 +65,18 @@ flowchart LR - HTTP server 不设置 WriteTimeout，避免杀死长时间流；上游只限制响应头等待时间，客户端取消会终止上游请求。 - 连接池按进程共享，默认保留 4096 个空闲连接、单 host 1024 个，可按实例并发和上游限制调整。 - 请求体默认最多 16 MiB，响应仅保留最多 64 KiB 用于非流式 usage 提取；正文直接传输。 -- 计量队列有界，慢消费者不会拖住推理请求。正式计费时不能仅靠“丢弃并打指标”，需要本地 WAL 或高可用消息系统。 +- 余额冻结和结算当前各需要一次 PostgreSQL 事务，以正确性优先。高吞吐阶段可把计费拆成独立服务并做账户分片，但不能用最终一致缓存替代权威账本事务。 + +## 管理面安全 + +bootstrap token 只映射为 `platform_admin`，用于首次签发控制台令牌和故障恢复。控制台令牌使用高熵随机值，数据库仅保存 SHA-256 摘要；平台角色没有 `tenant_id`，租户角色必须绑定一个 tenant。所有管理 API 在 handler 执行前校验 permission，租户过滤在 SQL 查询或资源所有权检查中完成，前端隐藏菜单不承担安全职责。 + +每个通过认证的管理请求都写入 `audit_logs`，包含 actor、角色、tenant、action、状态码、请求 ID、IP 和 User-Agent。审计写入失败不会回滚已成功的资源事务，但会输出结构化告警。 ## 建议的后续顺序 -1. 把管理 bootstrap token 替换为管理员账户、RBAC、审计事件，并与公网推理监听端口隔离。 -2. 增加价格表和账本：tenant、project、api_key、provider_credential、model、route、price_book。 -3. 持久 UsageEvent 和幂等账本，然后再实现余额预授权、扣费和退款。 -4. 主动健康检查、熔断、延迟 EWMA 和按成本路由。 -5. 规范化中间表示，实现可靠的跨协议转换与 Responses/Embeddings 等更多端点。 +1. 将管理监听端口与公网推理端口分离，并接入 OIDC/SSO、MFA 与短期会话。 +2. 增加价格版本、退款/冲正、Stripe dispute 处理和供应商日账单对账。 +3. 为不返回 usage 的上游增加可靠 token 计算器，并监控 `uncollected_micros`。 +4. 把 UsageEvent 做时间分区和归档，增加 CSV 导出与对账作业。 +5. 主动健康检查、熔断、延迟 EWMA 和按成本路由。 @@ -5,6 +5,7 @@ go 1.26 require ( github.com/jackc/pgx/v5 v5.10.0 github.com/redis/go-redis/v9 v9.21.0 + github.com/stripe/stripe-go/v86 v86.2.0 ) require ( @@ -1,6 +1,12 @@ +github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs= +github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c= +github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA= +github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= @@ -9,17 +15,30 @@ github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= +github.com/klauspost/cpuid/v2 v2.2.10 h1:tBs3QSyvjDyFTq3uoc/9xFpCuOsJQFNPiAhYdw2skhE= +github.com/klauspost/cpuid/v2 v2.2.10/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/redis/go-redis/v9 v9.21.0 h1:FPBE4hhbAke+TLmcY3WkpbDffJEomdqPn3HYiqAtL9E= github.com/redis/go-redis/v9 v9.21.0/go.mod h1:v/M13XI1PVCDcm01VtPFOADfZtHf8YW3baQf57KlIkA= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/stripe/stripe-go/v86 v86.2.0 h1:d6yW7JDnk03DEVGqct82j0rys8nPPCE3+RDPsIvUcrY= +github.com/stripe/stripe-go/v86 v86.2.0/go.mod h1:Co7QRXCKGNOPTugAdvjgRo+KcMtd9hxy+pZMN0yThsQ= +github.com/zeebo/xxh3 v1.1.0 h1:s7DLGDK45Dyfg7++yxI0khrfwq9661w9EN78eP/UZVs= +github.com/zeebo/xxh3 v1.1.0/go.mod h1:IisAie1LELR4xhVinxWS5+zf1lA4p0MW4T+w+W07F5s= go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE= go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0= golang.org/x/sync v0.17.0 h1:l60nONMj9l5drqw6jlhIELNv9I0A4OFgRsG9k2oT9Ug= golang.org/x/sync v0.17.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= +golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc= +golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/text v0.29.0 h1:1neNs90w9YzJ9BocxfsQNHKuAT4pkghyXc4nhZ6sJvk= golang.org/x/text v0.29.0/go.mod h1:7MhJOA9CD2qZyOKYazxdYMF85OwPdEr9jTtBpO7ydH4= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/internal/adminapi/api.go b/internal/adminapi/api.go index 1eb2a5d..7f5f8bd 100644 --- a/internal/adminapi/api.go +++ b/internal/adminapi/api.go @@ -1,15 +1,21 @@ package adminapi import ( + "context" + "crypto/rand" "crypto/subtle" + "encoding/hex" "encoding/json" "errors" "log/slog" + "net" "net/http" "strings" + "time" "aigw/internal/adminui" "aigw/internal/apierror" + "aigw/internal/billing" "aigw/internal/controlplane" "github.com/jackc/pgx/v5/pgconn" @@ -18,14 +24,33 @@ import ( type API struct { store *controlplane.Store manager *controlplane.Manager + billing *billing.Service token []byte logger *slog.Logger prefix string } +type actorKey struct{} +type auditWriter struct { + http.ResponseWriter + status int +} + +func (w *auditWriter) WriteHeader(status int) { + w.status = status + w.ResponseWriter.WriteHeader(status) +} +func (w *auditWriter) Write(body []byte) (int, error) { + if w.status == 0 { + w.status = http.StatusOK + } + return w.ResponseWriter.Write(body) +} + type Options struct { Store *controlplane.Store Manager *controlplane.Manager + Billing *billing.Service Token string Logger *slog.Logger Prefix string @@ -36,7 +61,7 @@ func New(options Options) *API { if prefix == "" { prefix = "/admin" } - return &API{store: options.Store, manager: options.Manager, token: []byte(options.Token), logger: options.Logger, prefix: prefix} + return &API{store: options.Store, manager: options.Manager, billing: options.Billing, token: []byte(options.Token), logger: options.Logger, prefix: prefix} } func (a *API) Handler() http.Handler { @@ -47,40 +72,119 @@ func (a *API) Handler() http.Handler { }) mux.Handle(a.prefix+"/", http.StripPrefix(a.prefix, adminui.Handler())) - mux.HandleFunc("GET "+apiPrefix+"/overview", a.withAuth(a.overview)) - mux.HandleFunc("GET "+apiPrefix+"/tenants", a.withAuth(a.listTenants)) - mux.HandleFunc("POST "+apiPrefix+"/tenants", a.withAuth(a.createTenant)) - mux.HandleFunc("GET "+apiPrefix+"/projects", a.withAuth(a.listProjects)) - mux.HandleFunc("POST "+apiPrefix+"/projects", a.withAuth(a.createProject)) - mux.HandleFunc("GET "+apiPrefix+"/keys", a.withAuth(a.listKeys)) - mux.HandleFunc("POST "+apiPrefix+"/keys", a.withAuth(a.createKey)) - mux.HandleFunc("POST "+apiPrefix+"/keys/{id}/revoke", a.withAuth(a.revokeKey)) - mux.HandleFunc("GET "+apiPrefix+"/providers", a.withAuth(a.listProviders)) - mux.HandleFunc("POST "+apiPrefix+"/providers", a.withAuth(a.createProvider)) - mux.HandleFunc("POST "+apiPrefix+"/providers/{id}/toggle", a.withAuth(a.toggleProvider)) - mux.HandleFunc("GET "+apiPrefix+"/models", a.withAuth(a.listModels)) - mux.HandleFunc("POST "+apiPrefix+"/models", a.withAuth(a.createModel)) - mux.HandleFunc("POST "+apiPrefix+"/models/{id}/toggle", a.withAuth(a.toggleModel)) - mux.HandleFunc("POST "+apiPrefix+"/reload", a.withAuth(a.reload)) + mux.HandleFunc("GET "+apiPrefix+"/overview", a.withAuth("overview.read", a.overview)) + mux.HandleFunc("GET "+apiPrefix+"/tenants", a.withAuth("tenants.read", a.listTenants)) + mux.HandleFunc("POST "+apiPrefix+"/tenants", a.withAuth("tenants.write", a.createTenant)) + mux.HandleFunc("GET "+apiPrefix+"/projects", a.withAuth("projects.read", a.listProjects)) + mux.HandleFunc("POST "+apiPrefix+"/projects", a.withAuth("projects.write", a.createProject)) + mux.HandleFunc("GET "+apiPrefix+"/keys", a.withAuth("keys.read", a.listKeys)) + mux.HandleFunc("POST "+apiPrefix+"/keys", a.withAuth("keys.write", a.createKey)) + mux.HandleFunc("POST "+apiPrefix+"/keys/{id}/revoke", a.withAuth("keys.write", a.revokeKey)) + mux.HandleFunc("GET "+apiPrefix+"/providers", a.withAuth("platform.read", a.listProviders)) + mux.HandleFunc("POST "+apiPrefix+"/providers", a.withAuth("platform.write", a.createProvider)) + mux.HandleFunc("POST "+apiPrefix+"/providers/{id}/toggle", a.withAuth("platform.write", a.toggleProvider)) + mux.HandleFunc("GET "+apiPrefix+"/models", a.withAuth("platform.read", a.listModels)) + mux.HandleFunc("POST "+apiPrefix+"/models", a.withAuth("platform.write", a.createModel)) + mux.HandleFunc("POST "+apiPrefix+"/models/{id}/toggle", a.withAuth("platform.write", a.toggleModel)) + mux.HandleFunc("POST "+apiPrefix+"/reload", a.withAuth("platform.write", a.reload)) + if a.billing != nil { + mux.HandleFunc("GET "+apiPrefix+"/billing/accounts", a.withAuth("billing.read", a.listBillingAccounts)) + mux.HandleFunc("GET "+apiPrefix+"/billing/ledger", a.withAuth("billing.read", a.listBillingLedger)) + mux.HandleFunc("POST "+apiPrefix+"/billing/adjustments", a.withAuth("billing.adjust", a.adjustBalance)) + mux.HandleFunc("POST "+apiPrefix+"/billing/checkout-sessions", a.withAuth("billing.topup", a.createCheckoutSession)) + } + mux.HandleFunc("GET "+apiPrefix+"/usage", a.withAuth("usage.read", a.listUsage)) + mux.HandleFunc("GET "+apiPrefix+"/usage/summary", a.withAuth("usage.read", a.usageSummary)) + mux.HandleFunc("GET "+apiPrefix+"/limits", a.withAuth("limits.read", a.listLimits)) + mux.HandleFunc("POST "+apiPrefix+"/limits/{project_id}", a.withAuth("limits.write", a.setLimit)) + mux.HandleFunc("GET "+apiPrefix+"/users", a.withAuth("users.read", a.listUsers)) + mux.HandleFunc("POST "+apiPrefix+"/users", a.withAuth("users.write", a.createUser)) + mux.HandleFunc("POST "+apiPrefix+"/users/{id}/revoke", a.withAuth("users.write", a.revokeUser)) + mux.HandleFunc("GET "+apiPrefix+"/audit", a.withAuth("audit.read", a.listAudit)) + mux.HandleFunc("GET "+apiPrefix+"/me", a.withAuth("overview.read", a.me)) return mux } -func (a *API) withAuth(next http.HandlerFunc) http.HandlerFunc { +func (a *API) withAuth(permission string, next http.HandlerFunc) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("X-AIGW-Request-ID") == "" { + r.Header.Set("X-AIGW-Request-ID", adminRequestID(r)) + } + w.Header().Set("X-AIGW-Request-ID", r.Header.Get("X-AIGW-Request-ID")) provided := strings.TrimSpace(r.Header.Get("X-Admin-Token")) if provided == "" { provided = bearerToken(r.Header.Get("Authorization")) } - if len(provided) == 0 || subtle.ConstantTimeCompare([]byte(provided), a.token) != 1 { + actor := controlplane.ConsoleActor{} + if len(provided) > 0 && len(a.token) > 0 && subtle.ConstantTimeCompare([]byte(provided), a.token) == 1 { + actor = controlplane.ConsoleActor{Role: controlplane.RolePlatformAdmin, DisplayName: "Bootstrap administrator", Bootstrap: true} + } else if provided != "" && a.store != nil { + var err error + actor, err = a.store.AuthenticateConsoleToken(r.Context(), provided) + if err != nil && !errors.Is(err, controlplane.ErrConsoleUnauthorized) { + a.logger.Error("console_authentication_failed", "error", err) + apierror.Write(w, apierror.Error{Status: http.StatusServiceUnavailable, Type: "control_plane_unavailable", Message: "Control plane authentication is temporarily unavailable"}, requestID(r)) + return + } + if err != nil { + actor = controlplane.ConsoleActor{} + } + } + if actor.Role == "" { apierror.Write(w, apierror.Error{Status: http.StatusUnauthorized, Type: "admin_unauthorized", Message: "Administrator authentication required"}, requestID(r)) return } - next(w, r) + if !actor.Can(permission) { + aw := &auditWriter{ResponseWriter: w} + apierror.Write(aw, apierror.Error{Status: http.StatusForbidden, Type: "admin_forbidden", Message: "You do not have permission for this operation"}, requestID(r)) + a.writeAudit(r, actor, permission, aw.status) + return + } + request := r.WithContext(context.WithValue(r.Context(), actorKey{}, actor)) + aw := &auditWriter{ResponseWriter: w} + next(aw, request) + if aw.status == 0 { + aw.status = http.StatusOK + } + a.writeAudit(request, actor, permission, aw.status) + } +} + +func (a *API) actor(r *http.Request) controlplane.ConsoleActor { + actor, _ := r.Context().Value(actorKey{}).(controlplane.ConsoleActor) + return actor +} + +func (a *API) writeAudit(r *http.Request, actor controlplane.ConsoleActor, action string, status int) { + if a.store == nil { + return + } + ip := r.RemoteAddr + if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil { + ip = host } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + if err := a.store.WriteAudit(ctx, controlplane.AuditInput{Actor: actor, RequestID: requestID(r), Method: r.Method, Path: r.URL.Path, Action: action, StatusCode: status, RemoteIP: ip, UserAgent: r.UserAgent()}); err != nil && a.logger != nil { + a.logger.Warn("admin_audit_write_failed", "error", err) + } +} + +func adminRequestID(r *http.Request) string { + if id := r.Header.Get("X-AIGW-Request-ID"); id != "" { + return id + } + var b [12]byte + if _, err := rand.Read(b[:]); err == nil { + return "adm_" + hex.EncodeToString(b[:]) + } + return "adm_unknown" } func (a *API) overview(w http.ResponseWriter, r *http.Request) { - result, err := a.store.Overview(r.Context()) + actor := a.actor(r) + tenantID := actor.TenantID + result, err := a.store.OverviewFor(r.Context(), tenantID) if err != nil { a.databaseError(w, r, err) return @@ -88,11 +192,67 @@ func (a *API) overview(w http.ResponseWriter, r *http.Request) { result.RuntimeGeneration = a.manager.Generation() result.RedisConfigured = a.manager.RedisConfigured() result.RedisConnected = a.manager.RedisConnected() + result.BillingEnabled = a.billing != nil + if a.billing != nil { + result.StripeEnabled = a.billing.StripeEnabled() + result.BillingCurrency = a.billing.Currency() + } + writeJSON(w, result) +} + +func (a *API) listBillingAccounts(w http.ResponseWriter, r *http.Request) { + result, err := a.billing.ListAccounts(r.Context(), a.actor(r).TenantID) + if err != nil { + a.billingError(w, r, err) + return + } writeJSON(w, result) } +func (a *API) listBillingLedger(w http.ResponseWriter, r *http.Request) { + tenantID := a.actor(r).TenantID + if tenantID == "" { + tenantID = r.URL.Query().Get("tenant_id") + } + result, err := a.billing.ListLedger(r.Context(), tenantID, 200) + if err != nil { + a.billingError(w, r, err) + return + } + writeJSON(w, result) +} + +func (a *API) adjustBalance(w http.ResponseWriter, r *http.Request) { + var input billing.AdjustmentInput + if !decodeBody(w, r, &input) { + return + } + result, err := a.billing.AdjustBalance(r.Context(), input) + if err != nil { + a.billingError(w, r, err) + return + } + writeStatusJSON(w, http.StatusCreated, result) +} + +func (a *API) createCheckoutSession(w http.ResponseWriter, r *http.Request) { + var input billing.CheckoutInput + if !decodeBody(w, r, &input) { + return + } + if tenantID := a.actor(r).TenantID; tenantID != "" { + input.TenantID = tenantID + } + result, err := a.billing.CreateCheckout(r.Context(), input) + if err != nil { + a.billingError(w, r, err) + return + } + writeStatusJSON(w, http.StatusCreated, result) +} + func (a *API) listTenants(w http.ResponseWriter, r *http.Request) { - result, err := a.store.ListTenants(r.Context()) + result, err := a.store.ListTenantsFor(r.Context(), a.actor(r).TenantID) if err != nil { a.databaseError(w, r, err) return @@ -117,7 +277,7 @@ func (a *API) createTenant(w http.ResponseWriter, r *http.Request) { } func (a *API) listProjects(w http.ResponseWriter, r *http.Request) { - result, err := a.store.ListProjects(r.Context()) + result, err := a.store.ListProjectsFor(r.Context(), a.actor(r).TenantID) if err != nil { a.databaseError(w, r, err) return @@ -130,6 +290,9 @@ func (a *API) createProject(w http.ResponseWriter, r *http.Request) { if !decodeBody(w, r, &input) { return } + if tenantID := a.actor(r).TenantID; tenantID != "" { + input.TenantID = tenantID + } result, generation, err := a.store.CreateProject(r.Context(), input) if err != nil { a.mutationError(w, r, err) @@ -142,7 +305,7 @@ func (a *API) createProject(w http.ResponseWriter, r *http.Request) { } func (a *API) listKeys(w http.ResponseWriter, r *http.Request) { - result, err := a.store.ListAPIKeys(r.Context()) + result, err := a.store.ListAPIKeysFor(r.Context(), a.actor(r).TenantID) if err != nil { a.databaseError(w, r, err) return @@ -155,6 +318,9 @@ func (a *API) createKey(w http.ResponseWriter, r *http.Request) { if !decodeBody(w, r, &input) { return } + if tenantID := a.actor(r).TenantID; tenantID != "" { + input.TenantID = tenantID + } result, generation, err := a.store.CreateAPIKey(r.Context(), input) if err != nil { a.mutationError(w, r, err) @@ -167,6 +333,10 @@ func (a *API) createKey(w http.ResponseWriter, r *http.Request) { } func (a *API) revokeKey(w http.ResponseWriter, r *http.Request) { + if err := a.requireResourceTenant(r, "api_key", r.PathValue("id")); err != nil { + a.scopeError(w, r) + return + } generation, err := a.store.RevokeAPIKey(r.Context(), r.PathValue("id")) if err != nil { a.mutationError(w, r, err) @@ -275,6 +445,127 @@ func (a *API) reload(w http.ResponseWriter, r *http.Request) { writeJSON(w, map[string]any{"generation": generation, "status": "reloaded"}) } +func (a *API) listUsage(w http.ResponseWriter, r *http.Request) { + query := controlplane.UsageQuery{TenantID: a.actor(r).TenantID, ProjectID: r.URL.Query().Get("project_id"), Model: r.URL.Query().Get("model"), Limit: 200} + result, err := a.store.ListUsage(r.Context(), query) + if err != nil { + a.databaseError(w, r, err) + return + } + writeJSON(w, result) +} + +func (a *API) usageSummary(w http.ResponseWriter, r *http.Request) { + result, err := a.store.UsageSummary(r.Context(), a.actor(r).TenantID, r.URL.Query().Get("project_id")) + if err != nil { + a.databaseError(w, r, err) + return + } + writeJSON(w, result) +} + +func (a *API) listLimits(w http.ResponseWriter, r *http.Request) { + result, err := a.store.ListProjectLimits(r.Context(), a.actor(r).TenantID) + if err != nil { + a.databaseError(w, r, err) + return + } + writeJSON(w, result) +} + +func (a *API) setLimit(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("project_id") + if err := a.requireResourceTenant(r, "project", id); err != nil { + a.scopeError(w, r) + return + } + var input controlplane.SetProjectLimitInput + if !decodeBody(w, r, &input) { + return + } + result, generation, err := a.store.SetProjectLimit(r.Context(), id, input) + if err != nil { + a.mutationError(w, r, err) + return + } + if !a.changed(w, r, generation, "project_limit", id) { + return + } + writeJSON(w, result) +} + +func (a *API) listUsers(w http.ResponseWriter, r *http.Request) { + result, err := a.store.ListConsoleUsers(r.Context(), a.actor(r).TenantID) + if err != nil { + a.databaseError(w, r, err) + return + } + writeJSON(w, result) +} + +func (a *API) createUser(w http.ResponseWriter, r *http.Request) { + var input controlplane.CreateConsoleUserInput + if !decodeBody(w, r, &input) { + return + } + if tenantID := a.actor(r).TenantID; tenantID != "" { + input.TenantID = tenantID + if strings.HasPrefix(input.Role, "platform_") { + a.scopeError(w, r) + return + } + } + result, err := a.store.CreateConsoleUser(r.Context(), input) + if err != nil { + a.mutationError(w, r, err) + return + } + writeStatusJSON(w, http.StatusCreated, result) +} + +func (a *API) revokeUser(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + if err := a.requireResourceTenant(r, "console_user", id); err != nil { + a.scopeError(w, r) + return + } + if err := a.store.RevokeConsoleUser(r.Context(), id, a.actor(r).TenantID); err != nil { + a.mutationError(w, r, err) + return + } + writeJSON(w, map[string]any{"status": "revoked", "id": id}) +} + +func (a *API) listAudit(w http.ResponseWriter, r *http.Request) { + result, err := a.store.ListAudit(r.Context(), a.actor(r).TenantID, 200) + if err != nil { + a.databaseError(w, r, err) + return + } + writeJSON(w, result) +} + +func (a *API) me(w http.ResponseWriter, r *http.Request) { + actor := a.actor(r) + writeJSON(w, map[string]any{"actor": actor, "permissions": actor.Permissions()}) +} + +func (a *API) requireResourceTenant(r *http.Request, resource, id string) error { + actor := a.actor(r) + if actor.IsPlatform() { + return nil + } + tenantID, err := a.store.ResourceTenantID(r.Context(), resource, id) + if err != nil || tenantID == "" || tenantID != actor.TenantID { + return errors.New("resource is outside tenant scope") + } + return nil +} + +func (a *API) scopeError(w http.ResponseWriter, r *http.Request) { + apierror.Write(w, apierror.Error{Status: http.StatusForbidden, Type: "tenant_scope", Message: "Resource is outside your tenant"}, requestID(r)) +} + func (a *API) changed(w http.ResponseWriter, r *http.Request, generation int64, resource, id string) bool { if err := a.manager.AfterMutation(r.Context(), generation, resource, id); err != nil { a.logger.Error("admin_control_plane_sync_failed", "resource", resource, "id", id, "error", err) @@ -314,6 +605,29 @@ func (a *API) mutationError(w http.ResponseWriter, r *http.Request, err error) { apierror.Write(w, apierror.Error{Status: status, Type: typeName, Message: message}, requestID(r)) } +func (a *API) billingError(w http.ResponseWriter, r *http.Request, err error) { + status := http.StatusInternalServerError + typeName := "billing_error" + message := "Billing operation failed" + switch { + case errors.Is(err, billing.ErrInvalidAmount): + status = http.StatusBadRequest + typeName = "invalid_amount" + message = "Amount is outside the configured bounds" + case errors.Is(err, billing.ErrInsufficientBalance): + status = http.StatusPaymentRequired + typeName = "insufficient_balance" + message = "Available balance is insufficient" + case errors.Is(err, billing.ErrStripeDisabled): + status = http.StatusServiceUnavailable + typeName = "stripe_disabled" + message = "Stripe top-ups are disabled" + default: + a.logger.Error("admin_billing_error", "error", err) + } + apierror.Write(w, apierror.Error{Status: status, Type: typeName, Message: message}, requestID(r)) +} + func decodeBody(w http.ResponseWriter, r *http.Request, destination any) bool { r.Body = http.MaxBytesReader(w, r.Body, 1<<20) defer r.Body.Close() @@ -348,5 +662,5 @@ func requestID(r *http.Request) string { if value := r.Header.Get("X-AIGW-Request-ID"); value != "" { return value } - return "admin" + return adminRequestID(r) } diff --git a/internal/adminui/assets/app.js b/internal/adminui/assets/app.js index f5ba048..6b8bf87 100644 --- a/internal/adminui/assets/app.js +++ b/internal/adminui/assets/app.js @@ -1,70 +1,131 @@ -const state = { token: sessionStorage.getItem('aigw_admin_token') || '', tenants: [], projects: [], providers: [], models: [], keys: [] }; +const state = { + token: sessionStorage.getItem('aigw_admin_token') || '', actor: {}, permissions: new Set(), overview: {}, + tenants: [], projects: [], keys: [], providers: [], models: [], billingAccounts: [], ledger: [], + usage: [], usageSummary: [], limits: [], users: [], audit: [] +}; const $ = (selector) => document.querySelector(selector); const $$ = (selector) => [...document.querySelectorAll(selector)]; +const can = (permission) => state.permissions.has(permission); -function esc(value) { - return String(value ?? '').replace(/[&<>'"]/g, (char) => ({ '&':'&', '<':'<', '>':'>', "'":''', '"':'"' }[char])); -} +function esc(value) { return String(value ?? '').replace(/[&<>'"]/g, (char) => ({ '&':'&', '<':'<', '>':'>', "'":''', '"':'"' }[char])); } function date(value) { return value ? new Date(value).toLocaleString() : '—'; } -function toast(message, error = false) { - const node = $('#toast'); node.textContent = message; node.className = `toast visible ${error ? 'error' : ''}`; - setTimeout(() => { node.className = 'toast'; }, 3200); -} +function shortID(value) { const text = String(value || ''); return text ? `${text.slice(0, 10)}${text.length > 10 ? '…' : ''}` : '—'; } +function percent(part, total) { return total ? `${Math.round((part / total) * 100)}%` : '—'; } +function toast(message, error = false) { const node = $('#toast'); node.textContent = message; node.className = `toast visible ${error ? 'error' : ''}`; setTimeout(() => { node.className = 'toast'; }, 3200); } async function api(path, options = {}) { - const response = await fetch(`./api${path}`, { ...options, headers: { 'Content-Type': 'application/json', 'Authorization': `Bearer ${state.token}`, ...(options.headers || {}) } }); + const response = await fetch(`./api${path}`, { ...options, headers: { 'Content-Type':'application/json', 'Authorization':`Bearer ${state.token}`, ...(options.headers || {}) } }); const payload = await response.json().catch(() => ({})); if (!response.ok) throw new Error(payload?.error?.message || `Request failed (${response.status})`); return payload; } function setConnected(connected) { - $('#connection-state').textContent = connected ? 'Connected' : 'Offline'; + $('#connection-state').textContent = connected ? state.actor.role?.replaceAll('_', ' ') || 'Connected' : 'Offline'; $('#connection-state').className = `state ${connected ? 'online' : ''}`; + $('#actor-label').textContent = connected ? state.actor.display_name || state.actor.email || '' : ''; } function formJSON(form) { return Object.fromEntries(new FormData(form).entries()); } function selectOptions(items, valueKey, labelKey, empty = 'Select…') { return `<option value="">${empty}</option>${items.map(item => `<option value="${esc(item[valueKey])}">${esc(item[labelKey])}</option>`).join('')}`; } +function decimalToScaled(value, digits) { + const match = String(value).trim().match(/^(-?)(\d+)(?:\.(\d+))?$/); if (!match) throw new Error('Enter a valid decimal amount'); + const fraction = match[3] || ''; if (fraction.length > digits && /[1-9]/.test(fraction.slice(digits))) throw new Error(`Use at most ${digits} decimal places`); + const scale = 10n ** BigInt(digits); const absolute = BigInt(match[2]) * scale + BigInt((fraction.slice(0, digits) + '0'.repeat(digits)).slice(0, digits) || '0'); + const result = Number(match[1] ? -absolute : absolute); if (!Number.isSafeInteger(result)) throw new Error('Amount is too large'); return result; +} +function scaledToDecimal(value, digits) { const number = BigInt(value || 0); const whole = number / (10n ** BigInt(digits)); const fraction = String(number % (10n ** BigInt(digits))).padStart(digits, '0').replace(/0+$/, ''); return fraction ? `${whole}.${fraction}` : String(whole); } +function currencyDigits(currency) { return ['bif','clp','djf','gnf','jpy','kmf','krw','mga','pyg','rwf','ugx','vnd','vuv','xaf','xof','xpf'].includes(currency) ? 0 : ['bhd','jod','kwd','omr','tnd'].includes(currency) ? 3 : 2; } +function money(micros, currency = state.overview.billing_currency || 'usd') { return new Intl.NumberFormat(undefined, { style:'currency', currency:currency.toUpperCase(), minimumFractionDigits:2, maximumFractionDigits:6 }).format(Number(micros || 0) / 1_000_000); } +function integer(value) { return new Intl.NumberFormat().format(Number(value || 0)); } +function emptyRow(span) { return `<tr><td colspan="${span}" class="empty">No records yet</td></tr>`; } +function showSecret(title, value) { $('#secret-title').textContent = title; $('#created-secret').textContent = value; $('#secret-dialog').showModal(); } +async function permitted(permission, path) { if (!can(permission)) return []; return api(path); } async function loadAll() { if (!state.token) { setConnected(false); return; } try { - [state.tenants, state.projects, state.providers, state.models, state.keys] = await Promise.all([ - api('/tenants'), api('/projects'), api('/providers'), api('/models'), api('/keys') + const session = await api('/me'); state.actor = session.actor || {}; state.permissions = new Set(session.permissions || []); + state.overview = await api('/overview'); + const results = await Promise.all([ + permitted('tenants.read','/tenants'), permitted('projects.read','/projects'), permitted('keys.read','/keys'), + permitted('platform.read','/providers'), permitted('platform.read','/models'), permitted('usage.read','/usage'), + permitted('usage.read','/usage/summary'), permitted('limits.read','/limits'), permitted('users.read','/users'), + permitted('audit.read','/audit'), state.overview.billing_enabled ? permitted('billing.read','/billing/accounts') : [], + state.overview.billing_enabled ? permitted('billing.read','/billing/ledger') : [] ]); - const overview = await api('/overview'); - renderOverview(overview); renderTenants(); renderProjects(); renderKeys(); renderProviders(); renderModels(); renderRouteEditor(); setConnected(true); + [state.tenants,state.projects,state.keys,state.providers,state.models,state.usage,state.usageSummary,state.limits,state.users,state.audit,state.billingAccounts,state.ledger] = results; + renderAll(); setConnected(true); } catch (error) { setConnected(false); toast(error.message, true); } } -function renderOverview(data) { - const propagation = !data.redis_configured ? 'PG polling' : data.redis_connected ? 'Redis live' : 'PG fallback'; - const items = [['Tenants', data.tenants, 'active accounts'], ['Projects', data.projects, 'workspaces'], ['API keys', data.api_keys, 'active credentials'], ['Providers', data.providers, 'enabled upstreams'], ['Models', data.models, 'public routes'], ['Runtime', data.runtime_generation, `PG generation ${data.generation}`], ['Propagation', propagation, data.redis_configured ? 'automatic recovery' : 'Redis not configured']]; - $('#metrics').innerHTML = items.map(([label, value, sub]) => `<article class="metric"><span>${label}</span><strong>${esc(value)}</strong><small>${sub}</small></article>`).join(''); + +function applyPermissions() { + $$('[data-permission]').forEach(node => node.classList.toggle('hidden', !can(node.dataset.permission))); + $('#billing-tab').classList.toggle('hidden', !state.overview.billing_enabled || !can('billing.read')); + $('#topup-form').classList.toggle('hidden', !state.overview.stripe_enabled || !can('billing.topup')); + const active = $('.tab.active'); if (active?.classList.contains('hidden')) $('.tab[data-section="overview"]').click(); } -function renderTenants() { $('#tenants-body').innerHTML = state.tenants.map(item => `<tr><td><strong>${esc(item.name)}</strong></td><td><code>${esc(item.slug)}</code></td><td><span class="badge ${item.status}">${esc(item.status)}</span></td><td>${date(item.created_at)}</td></tr>`).join('') || emptyRow(4); $('#project-tenant').innerHTML = selectOptions(state.tenants, 'id', 'name'); $('#key-tenant').innerHTML = selectOptions(state.tenants, 'id', 'name'); } -function renderProjects() { $('#projects-body').innerHTML = state.projects.map(item => `<tr><td><strong>${esc(item.name)}</strong></td><td><code>${esc(item.tenant_id).slice(0, 8)}…</code></td><td>${esc(item.slug)}</td><td><span class="badge ${item.status}">${esc(item.status)}</span></td></tr>`).join('') || emptyRow(4); renderKeyProjects(); } -function renderKeyProjects() { const tenant = $('#key-tenant').value; const projects = state.projects.filter(item => !tenant || item.tenant_id === tenant); $('#key-project').innerHTML = selectOptions(projects, 'id', 'name'); } -function renderKeys() { $('#keys-body').innerHTML = state.keys.map(item => `<tr><td><strong>${esc(item.name)}</strong></td><td><code>${esc(item.key_prefix)}</code></td><td><code>${esc(item.project_id).slice(0, 8)}…</code></td><td>${(item.scopes || []).map(scope => `<span class="tag">${esc(scope)}</span>`).join('')}</td><td><span class="badge ${item.status}">${esc(item.status)}</span></td><td>${item.status === 'active' ? `<button class="text-button danger" data-revoke-key="${esc(item.id)}">Revoke</button>` : ''}</td></tr>`).join('') || emptyRow(6); } -function renderProviders() { $('#providers-body').innerHTML = state.providers.map(item => `<tr><td><strong>${esc(item.name)}</strong></td><td><span class="tag">${esc(item.protocol)}</span></td><td class="truncate">${esc(item.base_url)}</td><td>${esc(item.route_count)}</td><td><span class="badge ${item.enabled ? 'active' : 'suspended'}">${item.enabled ? 'enabled' : 'disabled'}</span></td><td><button class="text-button" data-toggle-provider="${esc(item.id)}" data-enabled="${!item.enabled}">${item.enabled ? 'Disable' : 'Enable'}</button></td></tr>`).join('') || emptyRow(6); } -function renderModels() { $('#models-body').innerHTML = state.models.map(item => `<tr><td><strong>${esc(item.public_id)}</strong></td><td>${esc(item.owned_by || '—')}</td><td><div class="route-list">${(item.routes || []).map(route => `<span>${esc(route.provider_name || route.provider_id).slice(0, 24)} → ${esc(route.upstream_model)} <em>p${route.priority} / w${route.weight}</em></span>`).join('')}</div></td><td><span class="badge ${item.enabled ? 'active' : 'suspended'}">${item.enabled ? 'enabled' : 'disabled'}</span></td><td><button class="text-button" data-toggle-model="${esc(item.id)}" data-enabled="${!item.enabled}">${item.enabled ? 'Disable' : 'Enable'}</button></td></tr>`).join('') || emptyRow(5); } -function renderRouteEditor() { const current = $('#route-editor'); if (!current.children.length) addRoute(); $$('.route-provider').forEach(select => { const selected = select.value; select.innerHTML = selectOptions(state.providers.filter(item => item.enabled), 'id', 'name', 'Provider…'); select.value = selected; }); } -function addRoute() { const wrapper = document.createElement('div'); wrapper.className = 'route-row'; wrapper.innerHTML = `<select class="route-provider" required></select><input class="route-upstream" required placeholder="Upstream model"><input class="route-priority" type="number" min="0" value="0" title="Priority"><input class="route-weight" type="number" min="1" max="100" value="100" title="Weight"><button class="icon-button remove-route" type="button" aria-label="Remove route">×</button>`; $('#route-editor').appendChild(wrapper); renderRouteEditor(); } -function emptyRow(span) { return `<tr><td colspan="${span}" class="empty">No records yet</td></tr>`; } +function renderAll() { + applyPermissions(); renderOverview(); renderTenants(); renderProjects(); renderKeys(); renderProviders(); renderModels(); renderBilling(); + renderUsage(); renderLimits(); renderUsers(); renderAudit(); renderRouteEditor(); +} +function renderOverview() { + const data = state.overview; const propagation = !data.redis_configured ? 'PG polling' : data.redis_connected ? 'Redis live' : 'PG fallback'; + const now = new Date(); const current = state.usageSummary.filter(item => { const period = new Date(item.period_start); return period.getUTCMonth() === now.getUTCMonth() && period.getUTCFullYear() === now.getUTCFullYear(); }); + const totals = current.reduce((acc,item) => { acc.requests += item.request_count; acc.success += item.successful_requests; acc.tokens += item.total_tokens; acc.cost += item.cost_micros; return acc; }, {requests:0,success:0,tokens:0,cost:0}); + const items = [['Requests', integer(totals.requests), 'this month'], ['Success rate', percent(totals.success,totals.requests), 'completed requests'], ['Tokens', integer(totals.tokens), 'this month'], ['Spend', money(totals.cost), 'metered cost'], ['Projects', data.projects, 'active workspaces'], ['Propagation', propagation, `generation ${data.runtime_generation}`]]; + $('#metrics').innerHTML = items.map(([label,value,sub]) => `<article class="metric"><span>${label}</span><strong>${esc(value)}</strong><small>${esc(sub)}</small></article>`).join(''); + $('#overview-usage-body').innerHTML = current.map(item => `<tr><td><strong>${esc(item.project_name)}</strong></td><td>${integer(item.request_count)}</td><td>${percent(item.successful_requests,item.request_count)}</td><td>${integer(item.total_tokens)}</td><td>${money(item.cost_micros)}</td><td class="${item.uncollected_micros ? 'money-negative':''}">${money(item.uncollected_micros)}</td></tr>`).join('') || emptyRow(6); +} +function renderTenants() { + $('#tenants-body').innerHTML = state.tenants.map(item => `<tr><td><strong>${esc(item.name)}</strong></td><td><code>${esc(item.slug)}</code></td><td><span class="badge ${item.status}">${esc(item.status)}</span></td><td>${date(item.created_at)}</td></tr>`).join('') || emptyRow(4); + ['project-tenant','key-tenant','topup-tenant','adjustment-tenant','user-tenant'].forEach(id => { const node=$(`#${id}`); if (node) node.innerHTML=selectOptions(state.tenants,'id','name', state.actor.tenant_id ? 'Current tenant' : 'Select tenant…'); }); + if (state.actor.tenant_id) ['project-tenant','key-tenant','topup-tenant','adjustment-tenant','user-tenant'].forEach(id => { const node=$(`#${id}`); if (node) node.value=state.actor.tenant_id; }); +} +function renderProjects() { $('#projects-body').innerHTML = state.projects.map(item => `<tr><td><strong>${esc(item.name)}</strong></td><td><code>${shortID(item.tenant_id)}</code></td><td>${esc(item.slug)}</td><td><span class="badge ${item.status}">${esc(item.status)}</span></td></tr>`).join('') || emptyRow(4); renderKeyProjects(); } +function renderKeyProjects() { const tenant = $('#key-tenant').value; const projects = state.projects.filter(item => !tenant || item.tenant_id === tenant); $('#key-project').innerHTML = selectOptions(projects,'id','name'); } +function renderKeys() { $('#keys-body').innerHTML = state.keys.map(item => `<tr><td><strong>${esc(item.name)}</strong></td><td><code>${esc(item.key_prefix)}</code></td><td><code>${shortID(item.project_id)}</code></td><td>${(item.scopes||[]).map(scope=>`<span class="tag">${esc(scope)}</span>`).join('')}</td><td><span class="badge ${item.status}">${esc(item.status)}</span></td><td>${item.status==='active'&&can('keys.write')?`<button class="text-button danger" data-revoke-key="${esc(item.id)}">Revoke</button>`:''}</td></tr>`).join('') || emptyRow(6); } +function renderProviders() { $('#providers-body').innerHTML = state.providers.map(item => `<tr><td><strong>${esc(item.name)}</strong></td><td><span class="tag">${esc(item.protocol)}</span></td><td class="truncate">${esc(item.base_url)}</td><td>${integer(item.route_count)}</td><td><span class="badge ${item.enabled?'active':'suspended'}">${item.enabled?'enabled':'disabled'}</span></td><td>${can('platform.write')?`<button class="text-button" data-toggle-provider="${esc(item.id)}" data-enabled="${!item.enabled}">${item.enabled?'Disable':'Enable'}</button>`:''}</td></tr>`).join('') || emptyRow(6); } +function renderModels() { $('#models-body').innerHTML = state.models.map(item => `<tr><td><strong>${esc(item.public_id)}</strong></td><td>${esc(item.owned_by||'—')}<small class="price-line">in ${money(item.input_price_micros_per_million)}/1M · out ${money(item.output_price_micros_per_million)}/1M</small></td><td><div class="route-list">${(item.routes||[]).map(route=>`<span>${esc(route.provider_name||route.provider_id).slice(0,24)} → ${esc(route.upstream_model)} <em>p${route.priority} / w${route.weight}</em></span>`).join('')}</div></td><td><span class="badge ${item.enabled?'active':'suspended'}">${item.enabled?'enabled':'disabled'}</span></td><td>${can('platform.write')?`<button class="text-button" data-toggle-model="${esc(item.id)}" data-enabled="${!item.enabled}">${item.enabled?'Disable':'Enable'}</button>`:''}</td></tr>`).join('') || emptyRow(5); } +function renderBilling() { + $('#billing-currency').textContent=(state.overview.billing_currency||'').toUpperCase(); + $('#billing-accounts-body').innerHTML=state.billingAccounts.map(item=>`<tr><td><strong>${esc(item.tenant_name)}</strong><br><code>${shortID(item.tenant_id)}</code></td><td>${money(item.balance_micros,item.currency)}</td><td>${money(item.reserved_micros,item.currency)}</td><td><strong>${money(item.available_micros,item.currency)}</strong></td><td>${date(item.updated_at)}</td></tr>`).join('')||emptyRow(5); + $('#billing-ledger-body').innerHTML=state.ledger.map(item=>`<tr><td>${date(item.created_at)}</td><td><code>${shortID(item.tenant_id)}</code></td><td><span class="tag">${esc(item.kind)}</span></td><td class="${item.amount_micros>=0?'money-positive':'money-negative'}">${money(item.amount_micros,item.currency)}</td><td>${money(item.balance_after_micros,item.currency)}</td><td title="${esc(item.description)}"><code>${shortID(item.source_id)}</code></td></tr>`).join('')||emptyRow(6); +} +function renderUsage() { + $('#usage-summary-body').innerHTML=state.usageSummary.map(item=>`<tr><td>${new Date(item.period_start).toLocaleDateString(undefined,{year:'numeric',month:'short'})}</td><td><strong>${esc(item.project_name)}</strong></td><td>${integer(item.request_count)}</td><td>${percent(item.successful_requests,item.request_count)}</td><td>${integer(item.input_tokens)}</td><td>${integer(item.output_tokens)}</td><td>${money(item.cost_micros)}</td></tr>`).join('')||emptyRow(7); + $('#usage-events-body').innerHTML=state.usage.map(item=>`<tr><td>${date(item.started_at)}</td><td><code title="${esc(item.request_id)}">${shortID(item.request_id)}</code></td><td>${esc(item.public_model)}</td><td><span class="badge ${item.success?'active':'suspended'}">${item.status_code}</span>${item.error_type?`<small class="error-label">${esc(item.error_type)}</small>`:''}</td><td>${integer(item.total_tokens)}</td><td>${money(item.cost_micros)}</td><td>${integer(item.duration_ms)} ms</td></tr>`).join('')||emptyRow(7); +} +function renderLimits() { + const existing=new Map(state.limits.map(item=>[item.project_id,item])); + $('#limits-body').innerHTML=state.projects.map(project=>{const item=existing.get(project.id)||{};return `<tr data-limit-project="${esc(project.id)}"><td><strong>${esc(project.name)}</strong><br><code>${shortID(project.id)}</code></td><td><input class="limit-rpm" type="number" min="0" value="${item.requests_per_minute||0}" ${can('limits.write')?'':'disabled'}></td><td><input class="limit-tpm" type="number" min="0" value="${item.tokens_per_minute||0}" ${can('limits.write')?'':'disabled'}></td><td><input class="limit-concurrency" type="number" min="0" value="${item.concurrent_requests||0}" ${can('limits.write')?'':'disabled'}></td><td><input class="limit-spend" inputmode="decimal" value="${scaledToDecimal(item.monthly_spend_micros||0,6)}" ${can('limits.write')?'':'disabled'}></td><td>${can('limits.write')?'<button class="button secondary save-limit">Save</button>':''}</td></tr>`}).join('')||emptyRow(6); +} +function renderUsers() { + const roles = state.actor.tenant_id ? [['tenant_admin','Tenant admin'],['tenant_billing','Billing'],['tenant_developer','Developer'],['tenant_viewer','Viewer']] : [['platform_admin','Platform admin'],['platform_viewer','Platform viewer'],['tenant_admin','Tenant admin'],['tenant_billing','Billing'],['tenant_developer','Developer'],['tenant_viewer','Viewer']]; + $('#user-role').innerHTML=roles.map(([value,label])=>`<option value="${value}">${label}</option>`).join(''); + $('#users-body').innerHTML=state.users.map(item=>`<tr><td><strong>${esc(item.display_name)}</strong><br><span class="muted">${esc(item.email)}</span></td><td><span class="tag">${esc(item.role.replaceAll('_',' '))}</span></td><td><code>${shortID(item.tenant_id)}</code></td><td><code>${esc(item.token_prefix)}</code></td><td>${date(item.last_used_at)}</td><td><span class="badge ${item.status}">${esc(item.status)}</span></td><td>${item.status==='active'&&can('users.write')?`<button class="text-button danger" data-revoke-user="${esc(item.id)}">Revoke</button>`:''}</td></tr>`).join('')||emptyRow(7); +} +function renderAudit() { $('#audit-body').innerHTML=state.audit.map(item=>`<tr><td>${date(item.created_at)}</td><td><span class="tag">${esc(item.actor_role.replaceAll('_',' '))}</span></td><td>${esc(item.action)}</td><td><code>${esc(item.method)}</code></td><td><span class="badge ${item.status_code<400?'active':'suspended'}">${item.status_code}</span></td><td><code>${shortID(item.request_id)}</code></td><td><code>${esc(item.remote_ip||'—')}</code></td></tr>`).join('')||emptyRow(7); } +function renderRouteEditor() { const current=$('#route-editor');if(!current.children.length&&can('platform.write'))addRoute();$$('.route-provider').forEach(select=>{const selected=select.value;select.innerHTML=selectOptions(state.providers.filter(item=>item.enabled),'id','name','Provider…');select.value=selected;}); } +function addRoute() { const wrapper=document.createElement('div');wrapper.className='route-row';wrapper.innerHTML='<select class="route-provider" required></select><input class="route-upstream" required placeholder="Upstream model"><input class="route-priority" type="number" min="0" value="0" title="Priority"><input class="route-weight" type="number" min="1" max="100" value="100" title="Weight"><button class="icon-button remove-route" type="button" aria-label="Remove route">×</button>';$('#route-editor').appendChild(wrapper);renderRouteEditor(); } -document.addEventListener('click', async (event) => { - const tab = event.target.closest('.tab'); if (tab) { $$('.tab').forEach(node => node.classList.toggle('active', node === tab)); $$('.section').forEach(node => node.classList.toggle('active', node.id === tab.dataset.section)); return; } - if (event.target.id === 'reload') { try { await api('/reload', { method: 'POST', body: '{}' }); await loadAll(); toast('Snapshot reloaded'); } catch (error) { toast(error.message, true); } } - if (event.target.id === 'add-route') addRoute(); - if (event.target.closest('.remove-route')) { event.target.closest('.route-row').remove(); } - const revoke = event.target.closest('[data-revoke-key]'); if (revoke && confirm('Revoke this API key?')) { try { await api(`/keys/${revoke.dataset.revokeKey}/revoke`, { method: 'POST', body: '{}' }); await loadAll(); toast('API key revoked'); } catch (error) { toast(error.message, true); } } - const provider = event.target.closest('[data-toggle-provider]'); if (provider) { try { await api(`/providers/${provider.dataset.toggleProvider}/toggle`, { method: 'POST', body: JSON.stringify({ enabled: provider.dataset.enabled === 'true' }) }); await loadAll(); toast('Provider updated'); } catch (error) { toast(error.message, true); } } - const model = event.target.closest('[data-toggle-model]'); if (model) { try { await api(`/models/${model.dataset.toggleModel}/toggle`, { method: 'POST', body: JSON.stringify({ enabled: model.dataset.enabled === 'true' }) }); await loadAll(); toast('Model updated'); } catch (error) { toast(error.message, true); } } +document.addEventListener('click',async(event)=>{ + const tab=event.target.closest('.tab');if(tab){$$('.tab').forEach(node=>node.classList.toggle('active',node===tab));$$('.section').forEach(node=>node.classList.toggle('active',node.id===tab.dataset.section));return;} + if(event.target.id==='reload'){try{await api('/reload',{method:'POST',body:'{}'});await loadAll();toast('Snapshot reloaded');}catch(error){toast(error.message,true);}} + if(event.target.id==='add-route')addRoute();if(event.target.closest('.remove-route'))event.target.closest('.route-row').remove(); + const revokeKey=event.target.closest('[data-revoke-key]');if(revokeKey&&confirm('Revoke this API key?')){try{await api(`/keys/${revokeKey.dataset.revokeKey}/revoke`,{method:'POST',body:'{}'});await loadAll();toast('API key revoked');}catch(error){toast(error.message,true);}} + const revokeUser=event.target.closest('[data-revoke-user]');if(revokeUser&&confirm('Revoke this console credential?')){try{await api(`/users/${revokeUser.dataset.revokeUser}/revoke`,{method:'POST',body:'{}'});await loadAll();toast('Console credential revoked');}catch(error){toast(error.message,true);}} + const provider=event.target.closest('[data-toggle-provider]');if(provider){try{await api(`/providers/${provider.dataset.toggleProvider}/toggle`,{method:'POST',body:JSON.stringify({enabled:provider.dataset.enabled==='true'})});await loadAll();toast('Provider updated');}catch(error){toast(error.message,true);}} + const model=event.target.closest('[data-toggle-model]');if(model){try{await api(`/models/${model.dataset.toggleModel}/toggle`,{method:'POST',body:JSON.stringify({enabled:model.dataset.enabled==='true'})});await loadAll();toast('Model updated');}catch(error){toast(error.message,true);}} + const save=event.target.closest('.save-limit');if(save){const row=save.closest('[data-limit-project]');try{await api(`/limits/${row.dataset.limitProject}`,{method:'POST',body:JSON.stringify({requests_per_minute:Number(row.querySelector('.limit-rpm').value),tokens_per_minute:Number(row.querySelector('.limit-tpm').value),concurrent_requests:Number(row.querySelector('.limit-concurrency').value),monthly_spend_micros:decimalToScaled(row.querySelector('.limit-spend').value,6)})});await loadAll();toast('Project limits updated');}catch(error){toast(error.message,true);}} }); -$('#key-tenant').addEventListener('change', renderKeyProjects); -$('#session-form').addEventListener('submit', async (event) => { event.preventDefault(); state.token = $('#admin-token').value.trim(); sessionStorage.setItem('aigw_admin_token', state.token); await loadAll(); }); -$('#tenant-form').addEventListener('submit', async (event) => { event.preventDefault(); try { await api('/tenants', { method: 'POST', body: JSON.stringify(formJSON(event.target)) }); event.target.reset(); await loadAll(); toast('Tenant created'); } catch (error) { toast(error.message, true); } }); -$('#project-form').addEventListener('submit', async (event) => { event.preventDefault(); try { await api('/projects', { method: 'POST', body: JSON.stringify(formJSON(event.target)) }); event.target.reset(); await loadAll(); toast('Project created'); } catch (error) { toast(error.message, true); } }); -$('#key-form').addEventListener('submit', async (event) => { event.preventDefault(); try { const data = formJSON(event.target); data.scopes = data.scopes.split(',').map(value => value.trim()).filter(Boolean); const result = await api('/keys', { method: 'POST', body: JSON.stringify(data) }); event.target.reset(); $('#created-secret').textContent = result.key; $('#secret-dialog').showModal(); await loadAll(); } catch (error) { toast(error.message, true); } }); -$('#provider-form').addEventListener('submit', async (event) => { event.preventDefault(); try { await api('/providers', { method: 'POST', body: JSON.stringify(formJSON(event.target)) }); event.target.reset(); await loadAll(); toast('Provider added'); } catch (error) { toast(error.message, true); } }); -$('#model-form').addEventListener('submit', async (event) => { event.preventDefault(); try { const data = formJSON(event.target); data.routes = $$('.route-row').map(row => ({ provider_id: row.querySelector('.route-provider').value, upstream_model: row.querySelector('.route-upstream').value, priority: Number(row.querySelector('.route-priority').value), weight: Number(row.querySelector('.route-weight').value) })); await api('/models', { method: 'POST', body: JSON.stringify(data) }); event.target.reset(); $('#route-editor').innerHTML = ''; renderRouteEditor(); await loadAll(); toast('Model created'); } catch (error) { toast(error.message, true); } }); -$('#close-dialog').addEventListener('click', () => $('#secret-dialog').close()); -$('#copy-secret').addEventListener('click', async () => { await navigator.clipboard.writeText($('#created-secret').textContent); toast('Key copied'); }); -$('#admin-token').value = state.token; -if (state.token) loadAll(); + +$('#key-tenant').addEventListener('change',renderKeyProjects); +$('#session-form').addEventListener('submit',async(event)=>{event.preventDefault();state.token=$('#admin-token').value.trim();sessionStorage.setItem('aigw_admin_token',state.token);await loadAll();}); +$('#tenant-form').addEventListener('submit',async(event)=>{event.preventDefault();try{await api('/tenants',{method:'POST',body:JSON.stringify(formJSON(event.target))});event.target.reset();await loadAll();toast('Tenant created');}catch(error){toast(error.message,true);}}); +$('#project-form').addEventListener('submit',async(event)=>{event.preventDefault();try{await api('/projects',{method:'POST',body:JSON.stringify(formJSON(event.target))});event.target.reset();await loadAll();toast('Project created');}catch(error){toast(error.message,true);}}); +$('#key-form').addEventListener('submit',async(event)=>{event.preventDefault();try{const data=formJSON(event.target);data.scopes=data.scopes.split(',').map(value=>value.trim()).filter(Boolean);const result=await api('/keys',{method:'POST',body:JSON.stringify(data)});event.target.reset();showSecret('API key created',result.key);await loadAll();}catch(error){toast(error.message,true);}}); +$('#provider-form').addEventListener('submit',async(event)=>{event.preventDefault();try{await api('/providers',{method:'POST',body:JSON.stringify(formJSON(event.target))});event.target.reset();await loadAll();toast('Provider added');}catch(error){toast(error.message,true);}}); +$('#model-form').addEventListener('submit',async(event)=>{event.preventDefault();try{const data=formJSON(event.target);data.input_price_micros_per_million=decimalToScaled(data.input_price,6);data.output_price_micros_per_million=decimalToScaled(data.output_price,6);data.cache_read_price_micros_per_million=decimalToScaled(data.cache_read_price,6);data.cache_write_price_micros_per_million=decimalToScaled(data.cache_write_price,6);delete data.input_price;delete data.output_price;delete data.cache_read_price;delete data.cache_write_price;data.routes=$$('.route-row').map(row=>({provider_id:row.querySelector('.route-provider').value,upstream_model:row.querySelector('.route-upstream').value,priority:Number(row.querySelector('.route-priority').value),weight:Number(row.querySelector('.route-weight').value)}));await api('/models',{method:'POST',body:JSON.stringify(data)});event.target.reset();$('#route-editor').innerHTML='';renderRouteEditor();await loadAll();toast('Model created');}catch(error){toast(error.message,true);}}); +$('#topup-form').addEventListener('submit',async(event)=>{event.preventDefault();try{const data=formJSON(event.target);const digits=currencyDigits(state.overview.billing_currency||'usd');const result=await api('/billing/checkout-sessions',{method:'POST',body:JSON.stringify({tenant_id:data.tenant_id,amount_minor:decimalToScaled(data.amount,digits)})});window.location.assign(result.url);}catch(error){toast(error.message,true);}}); +$('#adjustment-form').addEventListener('submit',async(event)=>{event.preventDefault();try{const data=formJSON(event.target);await api('/billing/adjustments',{method:'POST',body:JSON.stringify({tenant_id:data.tenant_id,amount_micros:decimalToScaled(data.amount,6),description:data.description})});event.target.reset();await loadAll();toast('Balance adjusted');}catch(error){toast(error.message,true);}}); +$('#user-form').addEventListener('submit',async(event)=>{event.preventDefault();try{const result=await api('/users',{method:'POST',body:JSON.stringify(formJSON(event.target))});event.target.reset();showSecret('Console token issued',result.token);await loadAll();}catch(error){toast(error.message,true);}}); +$('#close-dialog').addEventListener('click',()=>$('#secret-dialog').close());$('#copy-secret').addEventListener('click',async()=>{await navigator.clipboard.writeText($('#created-secret').textContent);toast('Credential copied');}); +$('#admin-token').value=state.token;applyPermissions();if(state.token)loadAll(); diff --git a/internal/adminui/assets/index.html b/internal/adminui/assets/index.html index 7869855..c299037 100644 --- a/internal/adminui/assets/index.html +++ b/internal/adminui/assets/index.html @@ -10,57 +10,98 @@ <body> <header class="topbar"> <div class="brand"><span class="brand-mark">A</span><div><strong>AIGW</strong><small>CONTROL PLANE</small></div></div> - <form class="session" id="session-form"><input id="admin-token" name="admin-token" type="password" placeholder="Admin token" autocomplete="current-password" aria-label="Admin token"><button type="submit">Connect</button><span id="connection-state" class="state">Offline</span></form> + <form class="session" id="session-form"><span id="actor-label" class="actor-label"></span><input class="visually-hidden" name="username" autocomplete="username" value="aigw-console" aria-hidden="true" tabindex="-1"><input id="admin-token" name="admin-token" type="password" placeholder="Console token" autocomplete="current-password" aria-label="Console token"><button type="submit">Connect</button><span id="connection-state" class="state">Offline</span></form> </header> <main class="shell"> <nav class="tabs" aria-label="Admin sections"> <button class="tab active" data-section="overview">Overview</button> - <button class="tab" data-section="tenants">Tenants</button> - <button class="tab" data-section="projects">Projects</button> - <button class="tab" data-section="keys">API keys</button> - <button class="tab" data-section="providers">Providers</button> - <button class="tab" data-section="models">Models & routes</button> + <button class="tab" data-section="usage" data-permission="usage.read">Usage</button> + <button class="tab" data-section="billing" data-permission="billing.read" id="billing-tab">Billing</button> + <button class="tab" data-section="projects" data-permission="projects.read">Projects</button> + <button class="tab" data-section="keys" data-permission="keys.read">API keys</button> + <button class="tab" data-section="limits" data-permission="limits.read">Limits</button> + <button class="tab" data-section="team" data-permission="users.read">Team</button> + <button class="tab" data-section="audit" data-permission="audit.read">Audit</button> + <button class="tab" data-section="tenants" data-permission="tenants.read">Tenants</button> + <button class="tab" data-section="providers" data-permission="platform.read">Providers</button> + <button class="tab" data-section="models" data-permission="platform.read">Models & routes</button> </nav> <section id="overview" class="section active"> - <div class="section-heading"><div><span class="eyebrow">OPERATIONS</span><h1>Control plane overview</h1></div><button class="button secondary" id="reload">Reload snapshot</button></div> + <div class="section-heading"><div><span class="eyebrow">OPERATIONS</span><h1>Account overview</h1></div><button class="button secondary" id="reload" data-permission="platform.write">Reload snapshot</button></div> <div class="metric-grid" id="metrics"></div> - <div class="panel note"><div class="note-icon">i</div><div><strong>Runtime snapshot</strong><p>Writes commit to PostgreSQL and apply locally first. Redis accelerates propagation when available; PostgreSQL polling keeps every gateway convergent.</p></div></div> + <div class="section-heading ledger-heading"><div><span class="eyebrow">CURRENT PERIOD</span><h2>Project usage</h2></div></div> + <div class="panel table-wrap"><table><thead><tr><th>Project</th><th>Requests</th><th>Success</th><th>Tokens</th><th>Cost</th><th>Uncollected</th></tr></thead><tbody id="overview-usage-body"></tbody></table></div> + </section> + + <section id="usage" class="section"> + <div class="section-heading"><div><span class="eyebrow">METERING</span><h1>Usage ledger</h1></div></div> + <div class="panel table-wrap"><table><thead><tr><th>Period</th><th>Project</th><th>Requests</th><th>Success</th><th>Input</th><th>Output</th><th>Cost</th></tr></thead><tbody id="usage-summary-body"></tbody></table></div> + <div class="section-heading ledger-heading"><div><span class="eyebrow">REQUESTS</span><h2>Recent events</h2></div></div> + <div class="panel table-wrap"><table><thead><tr><th>Time</th><th>Request</th><th>Model</th><th>Status</th><th>Tokens</th><th>Cost</th><th>Latency</th></tr></thead><tbody id="usage-events-body"></tbody></table></div> </section> <section id="tenants" class="section"> <div class="section-heading"><div><span class="eyebrow">IDENTITY</span><h1>Tenants</h1></div></div> - <form class="panel form-grid" id="tenant-form"><label>Slug<input name="slug" required pattern="[a-z0-9][a-z0-9-]{1,62}[a-z0-9]" placeholder="acme"></label><label>Name<input name="name" required placeholder="Acme Inc."></label><button class="button primary" type="submit">Create tenant</button></form> + <form class="panel form-grid" id="tenant-form" data-permission="tenants.write"><label>Slug<input name="slug" required pattern="[a-z0-9][a-z0-9-]{1,62}[a-z0-9]" placeholder="acme"></label><label>Name<input name="name" required placeholder="Acme Inc."></label><button class="button primary" type="submit">Create tenant</button></form> <div class="panel table-wrap"><table><thead><tr><th>Name</th><th>Slug</th><th>Status</th><th>Created</th></tr></thead><tbody id="tenants-body"></tbody></table></div> </section> <section id="projects" class="section"> <div class="section-heading"><div><span class="eyebrow">IDENTITY</span><h1>Projects</h1></div></div> - <form class="panel form-grid" id="project-form"><label>Tenant<select name="tenant_id" id="project-tenant" required></select></label><label>Slug<input name="slug" required placeholder="production"></label><label>Name<input name="name" required placeholder="Production API"></label><button class="button primary" type="submit">Create project</button></form> + <form class="panel form-grid" id="project-form" data-permission="projects.write"><label>Tenant<select name="tenant_id" id="project-tenant" required></select></label><label>Slug<input name="slug" required placeholder="production"></label><label>Name<input name="name" required placeholder="Production API"></label><button class="button primary" type="submit">Create project</button></form> <div class="panel table-wrap"><table><thead><tr><th>Name</th><th>Tenant</th><th>Slug</th><th>Status</th></tr></thead><tbody id="projects-body"></tbody></table></div> </section> <section id="keys" class="section"> <div class="section-heading"><div><span class="eyebrow">ACCESS</span><h1>API keys</h1></div></div> - <form class="panel form-grid" id="key-form"><label>Tenant<select name="tenant_id" id="key-tenant" required></select></label><label>Project<select name="project_id" id="key-project" required></select></label><label>Name<input name="name" required placeholder="CLI production key"></label><label>Scopes<input name="scopes" value="inference" placeholder="inference,admin"></label><button class="button primary" type="submit">Create key</button></form> + <form class="panel form-grid" id="key-form" data-permission="keys.write"><label>Tenant<select name="tenant_id" id="key-tenant" required></select></label><label>Project<select name="project_id" id="key-project" required></select></label><label>Name<input name="name" required placeholder="CLI production key"></label><label>Scopes<input name="scopes" value="inference" placeholder="inference"></label><button class="button primary" type="submit">Create key</button></form> <div class="panel warning"><strong>Key visibility</strong><span>The secret is shown only once after creation.</span></div> <div class="panel table-wrap"><table><thead><tr><th>Name</th><th>Prefix</th><th>Project</th><th>Scopes</th><th>Status</th><th></th></tr></thead><tbody id="keys-body"></tbody></table></div> </section> <section id="providers" class="section"> <div class="section-heading"><div><span class="eyebrow">UPSTREAMS</span><h1>Providers</h1></div></div> - <form class="panel form-grid" id="provider-form"><label>Name<input name="name" required placeholder="openai-primary"></label><label>Protocol<select name="protocol"><option value="openai">OpenAI</option><option value="anthropic">Anthropic</option></select></label><label>Base URL<input name="base_url" type="url" required placeholder="https://api.example.com/v1"></label><label>API key<input name="api_key" type="password" required autocomplete="new-password" placeholder="Stored encrypted"></label><button class="button primary" type="submit">Add provider</button></form> + <form class="panel form-grid" id="provider-form" data-permission="platform.write"><input class="visually-hidden" autocomplete="username" value="aigw-provider" aria-hidden="true" tabindex="-1"><label>Name<input name="name" required placeholder="openai-primary"></label><label>Protocol<select name="protocol"><option value="openai">OpenAI</option><option value="anthropic">Anthropic</option></select></label><label>Base URL<input name="base_url" type="url" required placeholder="https://api.example.com/v1"></label><label>API key<input name="api_key" type="password" required autocomplete="new-password" placeholder="Stored encrypted"></label><button class="button primary" type="submit">Add provider</button></form> <div class="panel table-wrap"><table><thead><tr><th>Name</th><th>Protocol</th><th>Base URL</th><th>Routes</th><th>Status</th><th></th></tr></thead><tbody id="providers-body"></tbody></table></div> </section> <section id="models" class="section"> <div class="section-heading"><div><span class="eyebrow">ROUTING</span><h1>Models & routes</h1></div></div> - <form class="panel form-grid" id="model-form"><label>Public model ID<input name="public_id" required placeholder="openai/gpt-4.1-mini"></label><label>Owned by<input name="owned_by" placeholder="openai"></label><div class="route-editor" id="route-editor"></div><button class="button subtle" type="button" id="add-route">Add route</button><button class="button primary" type="submit">Create model</button></form> + <form class="panel form-grid" id="model-form" data-permission="platform.write"><label>Public model ID<input name="public_id" required placeholder="openai/gpt-4.1-mini"></label><label>Owned by<input name="owned_by" placeholder="openai"></label><label>Input price / 1M<input name="input_price" inputmode="decimal" value="0" required></label><label>Output price / 1M<input name="output_price" inputmode="decimal" value="0" required></label><label>Cache read / 1M<input name="cache_read_price" inputmode="decimal" value="0" required></label><label>Cache write / 1M<input name="cache_write_price" inputmode="decimal" value="0" required></label><div class="route-editor" id="route-editor"></div><button class="button subtle" type="button" id="add-route">Add route</button><button class="button primary" type="submit">Create model</button></form> <div class="panel table-wrap"><table><thead><tr><th>Public ID</th><th>Owner</th><th>Routes</th><th>Status</th><th></th></tr></thead><tbody id="models-body"></tbody></table></div> </section> + + <section id="billing" class="section"> + <div class="section-heading"><div><span class="eyebrow">REVENUE</span><h1>Balances & ledger</h1></div><span class="currency-label" id="billing-currency"></span></div> + <div class="billing-actions"> + <form class="panel form-grid compact-form" id="topup-form" data-permission="billing.topup"><label>Tenant<select name="tenant_id" id="topup-tenant" required></select></label><label>Amount<input name="amount" inputmode="decimal" min="0" required placeholder="25.00"></label><button class="button primary" type="submit">Open Stripe Checkout</button></form> + <form class="panel form-grid compact-form" id="adjustment-form" data-permission="billing.adjust"><label>Tenant<select name="tenant_id" id="adjustment-tenant" required></select></label><label>Signed amount<input name="amount" inputmode="decimal" required placeholder="10.00 or -5.00"></label><label>Reference<input name="description" maxlength="240" placeholder="Support credit"></label><button class="button secondary" type="submit">Post adjustment</button></form> + </div> + <div class="panel table-wrap"><table><thead><tr><th>Tenant</th><th>Balance</th><th>Reserved</th><th>Available</th><th>Updated</th></tr></thead><tbody id="billing-accounts-body"></tbody></table></div> + <div class="section-heading ledger-heading"><div><span class="eyebrow">AUDIT</span><h2>Recent ledger entries</h2></div></div> + <div class="panel table-wrap"><table><thead><tr><th>Time</th><th>Tenant</th><th>Kind</th><th>Amount</th><th>Balance after</th><th>Reference</th></tr></thead><tbody id="billing-ledger-body"></tbody></table></div> + </section> + + <section id="limits" class="section"> + <div class="section-heading"><div><span class="eyebrow">GUARDRAILS</span><h1>Project limits</h1></div></div> + <div class="panel table-wrap"><table class="limits-table"><thead><tr><th>Project</th><th>Requests / min</th><th>Estimated tokens / min</th><th>Concurrent</th><th>Monthly spend</th><th></th></tr></thead><tbody id="limits-body"></tbody></table></div> + </section> + + <section id="team" class="section"> + <div class="section-heading"><div><span class="eyebrow">RBAC</span><h1>Console access</h1></div></div> + <form class="panel form-grid" id="user-form" data-permission="users.write"><label>Tenant<select name="tenant_id" id="user-tenant"></select></label><label>Email<input name="email" type="email" required autocomplete="email" placeholder="operator@example.com"></label><label>Display name<input name="display_name" required placeholder="Operations"></label><label>Role<select name="role" id="user-role" required></select></label><button class="button primary" type="submit">Issue console token</button></form> + <div class="panel warning"><strong>Token visibility</strong><span>The console token is shown only once after creation.</span></div> + <div class="panel table-wrap"><table><thead><tr><th>User</th><th>Role</th><th>Tenant</th><th>Token</th><th>Last used</th><th>Status</th><th></th></tr></thead><tbody id="users-body"></tbody></table></div> + </section> + + <section id="audit" class="section"> + <div class="section-heading"><div><span class="eyebrow">SECURITY</span><h1>Audit log</h1></div></div> + <div class="panel table-wrap"><table><thead><tr><th>Time</th><th>Actor</th><th>Action</th><th>Method</th><th>Status</th><th>Request</th><th>IP</th></tr></thead><tbody id="audit-body"></tbody></table></div> + </section> </main> <div id="toast" class="toast" role="status"></div> - <dialog id="secret-dialog"><div class="dialog-content"><div class="section-heading"><div><span class="eyebrow">ONE-TIME SECRET</span><h2>API key created</h2></div><button class="icon-button" id="close-dialog" aria-label="Close">×</button></div><p>Copy this key now. It will not be shown again.</p><code id="created-secret"></code><button class="button primary" id="copy-secret">Copy key</button></div></dialog> + <dialog id="secret-dialog"><div class="dialog-content"><div class="section-heading"><div><span class="eyebrow">ONE-TIME SECRET</span><h2 id="secret-title">Credential created</h2></div><button class="icon-button" id="close-dialog" aria-label="Close">×</button></div><p>Copy this credential now. It will not be shown again.</p><code id="created-secret"></code><button class="button primary" id="copy-secret">Copy credential</button></div></dialog> <script src="./app.js" defer></script> </body> </html> diff --git a/internal/adminui/assets/style.css b/internal/adminui/assets/style.css index dadb019..b158d5b 100644 --- a/internal/adminui/assets/style.css +++ b/internal/adminui/assets/style.css @@ -1,13 +1,16 @@ :root { --bg:#f3f5f7; --panel:#fff; --ink:#18212b; --muted:#71808e; --line:#dce3e8; --accent:#146c94; --accent-soft:#e5f2f7; --danger:#b4494d; --shadow:0 8px 24px rgba(29,47,61,.06); font-family:Inter,ui-sans-serif,system-ui,-apple-system,BlinkMacSystemFont,"Segoe UI",sans-serif; } * { box-sizing:border-box; } body { margin:0; color:var(--ink); background:var(--bg); font-size:14px; } button,input,select { font:inherit; } button { cursor:pointer; } -.topbar { height:72px; background:#102a3a; color:#fff; padding:0 32px; display:flex; align-items:center; justify-content:space-between; gap:24px; } .brand { display:flex; align-items:center; gap:11px; letter-spacing:0; } .brand-mark { width:32px; height:32px; display:grid; place-items:center; border:1px solid #8fd0df; color:#b8eef7; font-weight:800; } .brand strong { display:block; font-size:15px; } .brand small { color:#8ba9b9; font-size:9px; letter-spacing:0; } .session { display:flex; align-items:center; gap:8px; } .session input { width:220px; border:1px solid #3b5b6c; background:#18384b; color:#fff; padding:9px 11px; outline:none; } .session input::placeholder { color:#91acb9; } .session button { min-height:40px; border:1px solid #8fd0df; background:#b8eef7; color:#102a3a; padding:0 14px; font-weight:750; } .session button:hover { background:#d4f6fb; } .state { color:#9db0bb; font-size:12px; } .state.online { color:#86d5ad; } -.shell { width:min(1240px,calc(100% - 48px)); margin:28px auto 60px; } .tabs { display:flex; gap:4px; border-bottom:1px solid var(--line); margin-bottom:26px; overflow:auto; } .tab { white-space:nowrap; border:0; background:transparent; color:var(--muted); padding:12px 15px; border-bottom:2px solid transparent; } .tab.active { color:var(--accent); border-bottom-color:var(--accent); font-weight:700; } +.topbar { height:72px; background:#102a3a; color:#fff; padding:0 32px; display:flex; align-items:center; justify-content:space-between; gap:24px; } .brand { display:flex; align-items:center; gap:11px; letter-spacing:0; } .brand-mark { width:32px; height:32px; display:grid; place-items:center; border:1px solid #8fd0df; color:#b8eef7; font-weight:800; } .brand strong { display:block; font-size:15px; } .brand small { color:#8ba9b9; font-size:9px; letter-spacing:0; } .session { display:flex; align-items:center; gap:8px; } .session input { width:220px; border:1px solid #3b5b6c; background:#18384b; color:#fff; padding:9px 11px; outline:none; } .session input::placeholder { color:#91acb9; } .session button { min-height:40px; border:1px solid #8fd0df; background:#b8eef7; color:#102a3a; padding:0 14px; font-weight:750; } .session button:hover { background:#d4f6fb; } .state { color:#9db0bb; font-size:12px; text-transform:capitalize; } .state.online { color:#86d5ad; } .actor-label { max-width:190px; overflow:hidden; text-overflow:ellipsis; white-space:nowrap; color:#c6d7df; font-size:12px; } +.shell { width:min(1240px,calc(100% - 48px)); margin:28px auto 60px; } .tabs { display:flex; flex-wrap:wrap; gap:4px; border-bottom:1px solid var(--line); margin-bottom:26px; } .tab { white-space:nowrap; border:0; background:transparent; color:var(--muted); padding:12px 15px; border-bottom:2px solid transparent; } .tab.active { color:var(--accent); border-bottom-color:var(--accent); font-weight:700; } .section { display:none; } .section.active { display:block; } .section-heading { display:flex; justify-content:space-between; align-items:flex-end; gap:20px; margin-bottom:19px; } .eyebrow { color:var(--accent); font-size:10px; letter-spacing:0; font-weight:800; } h1 { font-size:28px; line-height:1.1; margin:7px 0 0; letter-spacing:0; } h2 { margin:4px 0 0; font-size:20px; } .metric-grid { display:grid; grid-template-columns:repeat(6,1fr); gap:12px; } .metric { background:var(--panel); border:1px solid var(--line); padding:18px; box-shadow:var(--shadow); } .metric span,.metric small { display:block; color:var(--muted); } .metric strong { display:block; font-size:28px; margin:12px 0 3px; font-weight:750; } .metric small { font-size:11px; } .panel { background:var(--panel); border:1px solid var(--line); box-shadow:var(--shadow); padding:20px; margin-bottom:16px; } .note,.warning { display:flex; align-items:flex-start; gap:12px; } .note-icon { flex:0 0 22px; height:22px; border:1px solid var(--accent); color:var(--accent); display:grid; place-items:center; font-weight:700; } .note p { margin:5px 0 0; color:var(--muted); } .warning { color:#6d5523; background:#fff9e9; border-color:#ead9a9; box-shadow:none; } .warning span { margin-left:8px; color:#887650; } .form-grid { display:grid; grid-template-columns:repeat(3,minmax(0,1fr)); align-items:end; gap:13px; } label { display:flex; flex-direction:column; gap:7px; color:var(--muted); font-size:12px; font-weight:650; } input,select { width:100%; border:1px solid var(--line); background:#fff; color:var(--ink); padding:10px 11px; min-height:40px; outline:none; } input:focus,select:focus { border-color:#69a9bf; box-shadow:0 0 0 3px var(--accent-soft); } .button { border:1px solid transparent; min-height:40px; padding:0 15px; font-weight:700; } .button.primary { color:#fff; background:var(--accent); } .button.primary:hover { background:#0d5879; } .button.secondary { color:var(--accent); background:var(--accent-soft); border-color:#c5e1ea; } .button.subtle { color:var(--accent); background:#fff; border-color:var(--line); grid-column:1; } .table-wrap { overflow:auto; padding:0; } table { width:100%; border-collapse:collapse; min-width:700px; } th,td { padding:14px 18px; text-align:left; border-bottom:1px solid var(--line); vertical-align:middle; } th { color:var(--muted); font-size:11px; font-weight:700; text-transform:uppercase; letter-spacing:0; background:#fbfcfd; } tbody tr:last-child td { border-bottom:0; } td { font-size:13px; } code { font-family:"SFMono-Regular",Consolas,monospace; font-size:12px; color:#486071; } .badge,.tag { display:inline-flex; align-items:center; padding:4px 7px; font-size:11px; line-height:1; } .badge { border:1px solid #d9e0e4; color:var(--muted); } .badge.active { color:#187151; background:#e9f7f0; border-color:#c7e9d9; } .badge.revoked,.badge.suspended { color:var(--danger); background:#fff0f0; border-color:#f0cccc; } .tag { color:#4c6572; background:#eef3f5; margin:2px 3px 2px 0; } .text-button { border:0; background:transparent; color:var(--accent); padding:5px 0; } .text-button.danger { color:var(--danger); } .empty { color:var(--muted); text-align:center; padding:32px; } .truncate { max-width:280px; overflow:hidden; text-overflow:ellipsis; white-space:nowrap; } .route-editor { grid-column:1/-1; display:flex; flex-direction:column; gap:8px; } .route-row { display:grid; grid-template-columns:1.4fr 1.4fr .6fr .6fr 34px; gap:8px; } .icon-button { border:1px solid var(--line); background:#fff; color:var(--muted); width:34px; height:34px; font-size:19px; } .route-list { display:flex; flex-direction:column; gap:3px; color:#486071; font-size:12px; } .route-list em { color:var(--muted); font-style:normal; margin-left:4px; } +.hidden { display:none !important; } .visually-hidden { position:absolute !important; width:1px !important; height:1px !important; padding:0 !important; margin:-1px !important; overflow:hidden !important; clip:rect(0,0,0,0) !important; white-space:nowrap !important; border:0 !important; } .billing-actions { display:grid; grid-template-columns:1fr 1fr; gap:16px; } .compact-form { grid-template-columns:1fr 1fr; } .compact-form .button { grid-column:1/-1; } .currency-label { color:var(--muted); font-size:12px; text-transform:uppercase; } .ledger-heading { margin-top:28px; } .money-positive { color:#187151; } .money-negative { color:var(--danger); } +.price-line { display:block; color:var(--muted); font-size:10px; margin-top:5px; white-space:nowrap; } +.muted,.error-label { display:block; color:var(--muted); font-size:11px; margin-top:4px; } .error-label { color:var(--danger); } .limits-table input { min-width:118px; padding:8px 9px; } .limits-table .button { min-height:36px; } input:disabled,select:disabled { background:#f5f7f8; color:#697985; cursor:not-allowed; } .toast { position:fixed; bottom:24px; right:24px; background:#102a3a; color:#fff; padding:12px 16px; opacity:0; transform:translateY(8px); pointer-events:none; transition:.2s; } .toast.visible { opacity:1; transform:none; } .toast.error { background:#8f3d42; } dialog { border:0; padding:0; width:min(460px,calc(100% - 32px)); box-shadow:0 18px 70px rgba(0,0,0,.22); } dialog::backdrop { background:rgba(16,42,58,.45); } .dialog-content { padding:24px; } .dialog-content p { color:var(--muted); } .dialog-content code { display:block; background:#f3f5f7; padding:15px; overflow:auto; color:var(--ink); margin:18px 0; } -@media (max-width:900px) { .metric-grid { grid-template-columns:repeat(3,1fr); } .form-grid { grid-template-columns:repeat(2,minmax(0,1fr)); } .form-grid .button.primary { grid-column:1/-1; } } -@media (max-width:620px) { .topbar { height:auto; padding:16px; align-items:flex-start; flex-direction:column; } .session { width:100%; } .session input { flex:1; width:auto; min-width:0; } .shell { width:calc(100% - 24px); margin-top:18px; } .metric-grid { grid-template-columns:repeat(2,1fr); } .form-grid { grid-template-columns:1fr; } .route-row { grid-template-columns:minmax(0,1fr) minmax(0,1fr) 34px; } .route-provider,.route-upstream { grid-column:1/-1; } h1 { font-size:24px; } } +@media (max-width:900px) { .metric-grid { grid-template-columns:repeat(3,1fr); } .form-grid { grid-template-columns:repeat(2,minmax(0,1fr)); } .form-grid .button.primary { grid-column:1/-1; } .billing-actions { grid-template-columns:1fr; } } +@media (max-width:620px) { .topbar { height:auto; padding:16px; align-items:flex-start; flex-direction:column; } .session { width:100%; display:grid; grid-template-columns:minmax(0,1fr) auto; } .session input { width:100%; min-width:0; } .session .actor-label,.session .state { grid-column:1/-1; } .shell { width:calc(100% - 24px); margin-top:18px; } .tabs { margin-bottom:20px; flex-wrap:nowrap; overflow:auto; } .metric-grid { grid-template-columns:repeat(2,minmax(0,1fr)); } .metric { padding:14px; } .metric strong { font-size:22px; overflow-wrap:anywhere; } .form-grid { grid-template-columns:1fr; } .route-row { grid-template-columns:minmax(0,1fr) minmax(0,1fr) 34px; } .route-provider,.route-upstream { grid-column:1/-1; } h1 { font-size:24px; } .section-heading { align-items:flex-start; } } diff --git a/internal/billing/ledger.go b/internal/billing/ledger.go new file mode 100644 index 0000000..28cbe4c --- /dev/null +++ b/internal/billing/ledger.go @@ -0,0 +1,171 @@ +package billing + +import ( + "context" + "errors" + "fmt" + "math" + "strings" + + "github.com/jackc/pgx/v5" +) + +func (s *Service) ListAccounts(ctx context.Context, tenantID string) ([]Account, error) { + query := ` + SELECT t.id::text, t.name, COALESCE(w.currency, $1), COALESCE(w.balance_micros, 0), + COALESCE(w.reserved_micros, 0), COALESCE(w.balance_micros - w.reserved_micros, 0), + COALESCE(w.updated_at, t.created_at) + FROM tenants t LEFT JOIN tenant_wallets w ON w.tenant_id = t.id` + args := []any{s.currency} + if strings.TrimSpace(tenantID) != "" { + query += ` WHERE t.id=$2` + args = append(args, tenantID) + } + query += ` ORDER BY t.name` + rows, err := s.db.Query(ctx, query, args...) + if err != nil { + return nil, fmt.Errorf("query billing accounts: %w", err) + } + defer rows.Close() + result := make([]Account, 0) + for rows.Next() { + var item Account + if err := rows.Scan(&item.TenantID, &item.TenantName, &item.Currency, &item.BalanceMicros, + &item.ReservedMicros, &item.AvailableMicros, &item.UpdatedAt); err != nil { + return nil, fmt.Errorf("scan billing account: %w", err) + } + result = append(result, item) + } + return result, rows.Err() +} + +func (s *Service) ListLedger(ctx context.Context, tenantID string, limit int) ([]LedgerEntry, error) { + if limit < 1 || limit > 500 { + limit = 200 + } + query := ` + SELECT id::text, tenant_id::text, COALESCE(project_id::text, ''), currency, amount_micros, + balance_after_micros, kind, source_type, source_id, description, created_at + FROM billing_ledger` + args := []any{} + if strings.TrimSpace(tenantID) != "" { + query += ` WHERE tenant_id = $1 ORDER BY created_at DESC LIMIT $2` + args = append(args, tenantID, limit) + } else { + query += ` ORDER BY created_at DESC LIMIT $1` + args = append(args, limit) + } + rows, err := s.db.Query(ctx, query, args...) + if err != nil { + return nil, fmt.Errorf("query billing ledger: %w", err) + } + defer rows.Close() + result := make([]LedgerEntry, 0) + for rows.Next() { + var item LedgerEntry + if err := rows.Scan(&item.ID, &item.TenantID, &item.ProjectID, &item.Currency, &item.AmountMicros, + &item.BalanceAfterMicros, &item.Kind, &item.SourceType, &item.SourceID, + &item.Description, &item.CreatedAt); err != nil { + return nil, fmt.Errorf("scan billing ledger: %w", err) + } + result = append(result, item) + } + return result, rows.Err() +} + +func (s *Service) AdjustBalance(ctx context.Context, input AdjustmentInput) (LedgerEntry, error) { + input.TenantID = strings.TrimSpace(input.TenantID) + input.Description = normalizeDescription(input.Description) + if input.TenantID == "" || input.AmountMicros == 0 { + return LedgerEntry{}, ErrInvalidAmount + } + tx, err := s.db.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.ReadCommitted}) + if err != nil { + return LedgerEntry{}, fmt.Errorf("begin balance adjustment: %w", err) + } + defer tx.Rollback(ctx) + if _, err := tx.Exec(ctx, ` + INSERT INTO tenant_wallets (tenant_id, currency) VALUES ($1, $2) + ON CONFLICT (tenant_id) DO NOTHING`, input.TenantID, s.currency); err != nil { + return LedgerEntry{}, fmt.Errorf("ensure adjustment wallet: %w", err) + } + var currency string + var balance, held int64 + if err := tx.QueryRow(ctx, `SELECT currency, balance_micros, reserved_micros FROM tenant_wallets WHERE tenant_id = $1 FOR UPDATE`, input.TenantID).Scan(¤cy, &balance, &held); err != nil { + return LedgerEntry{}, fmt.Errorf("lock adjustment wallet: %w", err) + } + if currency != s.currency { + return LedgerEntry{}, fmt.Errorf("tenant wallet currency %s does not match %s", currency, s.currency) + } + if input.AmountMicros > 0 && balance > math.MaxInt64-input.AmountMicros { + return LedgerEntry{}, ErrInvalidAmount + } + newBalance := balance + input.AmountMicros + if newBalance < held || newBalance < 0 { + return LedgerEntry{}, ErrInsufficientBalance + } + if _, err := tx.Exec(ctx, `UPDATE tenant_wallets SET balance_micros = $2, updated_at = now() WHERE tenant_id = $1`, input.TenantID, newBalance); err != nil { + return LedgerEntry{}, fmt.Errorf("apply balance adjustment: %w", err) + } + sourceID := "adj_" + randomHex(16) + var result LedgerEntry + err = tx.QueryRow(ctx, ` + INSERT INTO billing_ledger (tenant_id, currency, amount_micros, balance_after_micros, kind, source_type, source_id, description) + VALUES ($1,$2,$3,$4,'adjustment','admin',$5,$6) + RETURNING id::text, tenant_id::text, '', currency, amount_micros, balance_after_micros, + kind, source_type, source_id, description, created_at`, + input.TenantID, s.currency, input.AmountMicros, newBalance, sourceID, input.Description, + ).Scan(&result.ID, &result.TenantID, &result.ProjectID, &result.Currency, &result.AmountMicros, + &result.BalanceAfterMicros, &result.Kind, &result.SourceType, &result.SourceID, + &result.Description, &result.CreatedAt) + if err != nil { + return LedgerEntry{}, fmt.Errorf("write adjustment ledger entry: %w", err) + } + if err := tx.Commit(ctx); err != nil { + return LedgerEntry{}, fmt.Errorf("commit balance adjustment: %w", err) + } + return result, nil +} + +func (s *Service) createTopUpOrder(ctx context.Context, input CheckoutInput) (string, int64, error) { + if !s.stripeEnabled { + return "", 0, ErrStripeDisabled + } + if strings.TrimSpace(input.TenantID) == "" || input.AmountMinor < s.minTopUpMinor || input.AmountMinor > s.maxTopUpMinor { + return "", 0, ErrInvalidAmount + } + amountMicros, err := minorToMicros(s.currency, input.AmountMinor) + if err != nil { + return "", 0, err + } + var orderID string + err = s.db.QueryRow(ctx, ` + INSERT INTO topup_orders (tenant_id, amount_minor, amount_micros, currency) + VALUES ($1,$2,$3,$4) RETURNING id::text`, input.TenantID, input.AmountMinor, amountMicros, s.currency, + ).Scan(&orderID) + if err != nil { + return "", 0, fmt.Errorf("create top-up order: %w", err) + } + return orderID, amountMicros, nil +} + +func minorToMicros(currency string, amount int64) (int64, error) { + if amount <= 0 { + return 0, ErrInvalidAmount + } + factor := int64(10_000) + switch strings.ToLower(currency) { + case "bif", "clp", "djf", "gnf", "jpy", "kmf", "krw", "mga", "pyg", "rwf", "ugx", "vnd", "vuv", "xaf", "xof", "xpf": + factor = 1_000_000 + case "bhd", "jod", "kwd", "omr", "tnd": + factor = 1_000 + } + if amount > math.MaxInt64/factor { + return 0, ErrInvalidAmount + } + return amount * factor, nil +} + +func isNotFound(err error) bool { + return errors.Is(err, pgx.ErrNoRows) +} diff --git a/internal/billing/service.go b/internal/billing/service.go new file mode 100644 index 0000000..b87a0bd --- /dev/null +++ b/internal/billing/service.go @@ -0,0 +1,352 @@ +package billing + +import ( + "context" + "crypto/rand" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "math/big" + "strings" + "time" + + "aigw/internal/domain" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" +) + +const microsPerUnit = int64(1_000_000) + +type Service struct { + db *pgxpool.Pool + currency string + defaultMaxOutputTokens int64 + minTopUpMinor int64 + maxTopUpMinor int64 + stripeEnabled bool + stripeWebhookSecret string + stripeSuccessURL string + stripeCancelURL string + integrationIdentifier string + createStripeCheckout stripeCheckoutCreator +} + +func New(ctx context.Context, options Options) (*Service, error) { + db, err := pgxpool.New(ctx, options.DatabaseURL) + if err != nil { + return nil, fmt.Errorf("configure billing PostgreSQL: %w", err) + } + if err := db.Ping(ctx); err != nil { + db.Close() + return nil, fmt.Errorf("connect billing PostgreSQL: %w", err) + } + service := &Service{ + db: db, currency: options.Currency, defaultMaxOutputTokens: options.DefaultMaxOutputTokens, + minTopUpMinor: options.MinTopUpMinor, maxTopUpMinor: options.MaxTopUpMinor, + stripeEnabled: options.StripeEnabled, stripeWebhookSecret: options.StripeWebhookSecret, + stripeSuccessURL: options.StripeSuccessURL, stripeCancelURL: options.StripeCancelURL, + integrationIdentifier: "aigw_balance_" + randomLetters(8), + } + if options.StripeEnabled { + service.createStripeCheckout = newStripeCheckoutCreator(options.StripeAPIKey) + } + return service, nil +} + +func (s *Service) Close() { + s.db.Close() +} + +func (s *Service) StripeEnabled() bool { + return s.stripeEnabled +} + +func (s *Service) Currency() string { + return s.currency +} + +func (s *Service) Authorize(ctx context.Context, input Authorization) error { + if input.RequestID == "" || input.Principal.TenantID == "" || input.Principal.ProjectID == "" || input.Principal.KeyID == "" { + return errors.New("billing authorization identity is incomplete") + } + reserved, err := reservationCost(input.Model, input.Body, s.defaultMaxOutputTokens) + if err != nil { + return err + } + tx, err := s.db.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.ReadCommitted}) + if err != nil { + return fmt.Errorf("begin billing authorization: %w", err) + } + defer tx.Rollback(ctx) + if _, err := tx.Exec(ctx, ` + INSERT INTO tenant_wallets (tenant_id, currency) VALUES ($1, $2) + ON CONFLICT (tenant_id) DO NOTHING`, input.Principal.TenantID, s.currency); err != nil { + return fmt.Errorf("ensure tenant wallet: %w", err) + } + var currency string + var balance, held int64 + if err := tx.QueryRow(ctx, ` + SELECT currency, balance_micros, reserved_micros FROM tenant_wallets + WHERE tenant_id = $1 FOR UPDATE`, input.Principal.TenantID).Scan(¤cy, &balance, &held); err != nil { + return fmt.Errorf("lock tenant wallet: %w", err) + } + if currency != s.currency { + return fmt.Errorf("tenant wallet currency %s does not match billing currency %s", currency, s.currency) + } + if input.Policy.MonthlySpendMicros > 0 { + period := time.Date(time.Now().UTC().Year(), time.Now().UTC().Month(), 1, 0, 0, 0, 0, time.UTC) + nextPeriod := period.AddDate(0, 1, 0) + var used, pending int64 + if err := tx.QueryRow(ctx, `SELECT + COALESCE((SELECT cost_micros FROM usage_monthly_rollups WHERE project_id=$1 AND period_start=$2),0), + COALESCE((SELECT sum(reserved_micros) FROM billing_reservations WHERE project_id=$1 AND status='pending' AND created_at >= $2 AND created_at < $3),0)`, + input.Principal.ProjectID, period, nextPeriod).Scan(&used, &pending); err != nil { + return fmt.Errorf("read monthly spend quota: %w", err) + } + if reserved > input.Policy.MonthlySpendMicros || used > input.Policy.MonthlySpendMicros-reserved || pending > input.Policy.MonthlySpendMicros-used-reserved { + return ErrQuotaExceeded + } + } + if balance-held < reserved { + return ErrInsufficientBalance + } + if _, err := tx.Exec(ctx, ` + INSERT INTO billing_reservations ( + request_id, tenant_id, project_id, key_id, public_model, currency, reserved_micros, + input_price_micros_per_million, output_price_micros_per_million, + cache_read_price_micros_per_million, cache_write_price_micros_per_million) + VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11)`, + input.RequestID, input.Principal.TenantID, input.Principal.ProjectID, input.Principal.KeyID, + input.Model.ID, s.currency, reserved, input.Model.InputPriceMicrosPerMillion, + input.Model.OutputPriceMicrosPerMillion, input.Model.CacheReadPriceMicrosPerMillion, + input.Model.CacheWritePriceMicrosPerMillion); err != nil { + return fmt.Errorf("create billing reservation: %w", err) + } + if _, err := tx.Exec(ctx, ` + UPDATE tenant_wallets SET reserved_micros = reserved_micros + $2, updated_at = now() + WHERE tenant_id = $1`, input.Principal.TenantID, reserved); err != nil { + return fmt.Errorf("reserve tenant balance: %w", err) + } + if err := tx.Commit(ctx); err != nil { + return fmt.Errorf("commit billing authorization: %w", err) + } + return nil +} + +func (s *Service) Settle(ctx context.Context, event domain.UsageEvent) error { + tx, err := s.db.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.ReadCommitted}) + if err != nil { + return fmt.Errorf("begin usage settlement: %w", err) + } + defer tx.Rollback(ctx) + + var tenantID, projectID, keyID, modelID, currency, status string + var reserved, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice int64 + err = tx.QueryRow(ctx, ` + SELECT tenant_id::text, project_id::text, key_id::text, public_model, currency, reserved_micros, status, + input_price_micros_per_million, output_price_micros_per_million, + cache_read_price_micros_per_million, cache_write_price_micros_per_million + FROM billing_reservations WHERE request_id = $1 FOR UPDATE`, event.RequestID, + ).Scan(&tenantID, &projectID, &keyID, &modelID, ¤cy, &reserved, &status, + &inputPrice, &outputPrice, &cacheReadPrice, &cacheWritePrice) + if err != nil { + return fmt.Errorf("load billing reservation: %w", err) + } + if status != "pending" { + return tx.Commit(ctx) + } + + actualCost := int64(0) + if event.StatusCode >= 200 && event.StatusCode < 300 { + actualCost, err = usageCost(event.Usage, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice) + if err != nil { + return err + } + } + var balance, held int64 + if err := tx.QueryRow(ctx, `SELECT balance_micros, reserved_micros FROM tenant_wallets WHERE tenant_id = $1 FOR UPDATE`, tenantID).Scan(&balance, &held); err != nil { + return fmt.Errorf("lock wallet for settlement: %w", err) + } + charged, err := collectibleCharge(actualCost, balance, held, reserved) + if err != nil { + return err + } + uncollected := actualCost - charged + newBalance := balance - charged + if _, err := tx.Exec(ctx, ` + UPDATE tenant_wallets SET balance_micros = $2, reserved_micros = reserved_micros - $3, updated_at = now() + WHERE tenant_id = $1`, tenantID, newBalance, reserved); err != nil { + return fmt.Errorf("settle tenant wallet: %w", err) + } + reservationStatus := "released" + if actualCost > 0 { + reservationStatus = "settled" + } + if _, err := tx.Exec(ctx, ` + UPDATE billing_reservations + SET status = $2, actual_cost_micros = $3, charged_micros = $4, uncollected_micros = $5, settled_at = now() + WHERE request_id = $1`, event.RequestID, reservationStatus, actualCost, charged, uncollected); err != nil { + return fmt.Errorf("update billing reservation: %w", err) + } + usageAlreadyRecorded := false + if err := tx.QueryRow(ctx, `SELECT true FROM usage_events WHERE request_id=$1 FOR UPDATE`, event.RequestID).Scan(&usageAlreadyRecorded); err != nil && !errors.Is(err, pgx.ErrNoRows) { + return fmt.Errorf("lock existing usage event: %w", err) + } + if usageAlreadyRecorded { + if _, err := tx.Exec(ctx, `UPDATE usage_events SET cost_micros=$2, charged_micros=$3, uncollected_micros=$4 WHERE request_id=$1`, + event.RequestID, actualCost, charged, uncollected); err != nil { + return fmt.Errorf("apply usage charge: %w", err) + } + } else if _, err := tx.Exec(ctx, ` + INSERT INTO usage_events ( + request_id, tenant_id, project_id, key_id, public_model, provider_id, upstream_model, + protocol, stream, status_code, success, error_type, attempts, started_at, duration_ms, + input_tokens, output_tokens, total_tokens, cache_creation_input_tokens, cache_read_input_tokens, + cost_micros, charged_micros, uncollected_micros) + VALUES ($1,$2,$3,$4,$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21,$22,$23) + ON CONFLICT (request_id) DO NOTHING`, + event.RequestID, tenantID, projectID, keyID, modelID, event.ProviderID, event.UpstreamModel, + string(event.Protocol), event.Stream, event.StatusCode, event.Success, event.ErrorType, event.Attempts, + event.StartedAt, event.DurationMS, event.Usage.InputTokens, event.Usage.OutputTokens, + event.Usage.TotalTokens, event.Usage.CacheCreationInputTokens, event.Usage.CacheReadInputTokens, + actualCost, charged, uncollected); err != nil { + return fmt.Errorf("persist usage event: %w", err) + } + if charged > 0 { + if _, err := tx.Exec(ctx, ` + INSERT INTO billing_ledger (tenant_id, project_id, currency, amount_micros, balance_after_micros, kind, source_type, source_id, description) + VALUES ($1,$2,$3,$4,$5,'usage','request',$6,$7) + ON CONFLICT (source_type, source_id) DO NOTHING`, tenantID, projectID, currency, -charged, newBalance, event.RequestID, modelID); err != nil { + return fmt.Errorf("write usage ledger entry: %w", err) + } + } + period := time.Date(event.StartedAt.UTC().Year(), event.StartedAt.UTC().Month(), 1, 0, 0, 0, 0, time.UTC) + if usageAlreadyRecorded { + if _, err := tx.Exec(ctx, `UPDATE usage_monthly_rollups SET cost_micros=cost_micros+$3, + charged_micros=charged_micros+$4, uncollected_micros=uncollected_micros+$5, updated_at=now() + WHERE project_id=$1 AND period_start=$2`, projectID, period, actualCost, charged, uncollected); err != nil { + return fmt.Errorf("apply usage rollup charge: %w", err) + } + } else if _, err := tx.Exec(ctx, ` + INSERT INTO usage_monthly_rollups + (period_start, tenant_id, project_id, request_count, successful_requests, input_tokens, output_tokens, total_tokens, cost_micros, charged_micros, uncollected_micros) + VALUES ($1,$2,$3,1,$4,$5,$6,$7,$8,$9,$10) + ON CONFLICT (project_id, period_start) DO UPDATE SET request_count=usage_monthly_rollups.request_count+1, + successful_requests=usage_monthly_rollups.successful_requests+EXCLUDED.successful_requests, + input_tokens=usage_monthly_rollups.input_tokens+EXCLUDED.input_tokens, + output_tokens=usage_monthly_rollups.output_tokens+EXCLUDED.output_tokens, + total_tokens=usage_monthly_rollups.total_tokens+EXCLUDED.total_tokens, + cost_micros=usage_monthly_rollups.cost_micros+EXCLUDED.cost_micros, + charged_micros=usage_monthly_rollups.charged_micros+EXCLUDED.charged_micros, + uncollected_micros=usage_monthly_rollups.uncollected_micros+EXCLUDED.uncollected_micros, + updated_at=now()`, period, tenantID, projectID, boolToInt(event.Success), event.Usage.InputTokens, + event.Usage.OutputTokens, event.Usage.TotalTokens, actualCost, charged, uncollected); err != nil { + return fmt.Errorf("update usage monthly rollup: %w", err) + } + if err := tx.Commit(ctx); err != nil { + return fmt.Errorf("commit usage settlement: %w", err) + } + return nil +} + +func boolToInt(value bool) int { + if value { + return 1 + } + return 0 +} + +func reservationCost(model domain.Model, body []byte, defaultMaxOutput int64) (int64, error) { + maxOutput := defaultMaxOutput + var limits struct { + MaxTokens int64 `json:"max_tokens"` + MaxCompletionTokens int64 `json:"max_completion_tokens"` + MaxOutputTokens int64 `json:"max_output_tokens"` + } + if json.Unmarshal(body, &limits) == nil { + explicitMax := int64(0) + for _, value := range []int64{limits.MaxTokens, limits.MaxCompletionTokens, limits.MaxOutputTokens} { + if value > explicitMax { + explicitMax = value + } + } + if explicitMax > 0 { + maxOutput = explicitMax + } + } + cacheReservePrice := model.CacheReadPriceMicrosPerMillion + if model.CacheWritePriceMicrosPerMillion > cacheReservePrice { + cacheReservePrice = model.CacheWritePriceMicrosPerMillion + } + return calculateCost(int64(len(body)), maxOutput, 0, int64(len(body)), + model.InputPriceMicrosPerMillion, model.OutputPriceMicrosPerMillion, + 0, cacheReservePrice) +} + +func usageCost(usage domain.Usage, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice int64) (int64, error) { + return calculateCost(usage.InputTokens, usage.OutputTokens, usage.CacheReadInputTokens, + usage.CacheCreationInputTokens, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice) +} + +func calculateCost(input, output, cacheRead, cacheWrite, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice int64) (int64, error) { + values := []int64{input, output, cacheRead, cacheWrite, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice} + for _, value := range values { + if value < 0 { + return 0, errors.New("billing values cannot be negative") + } + } + total := new(big.Int) + for _, pair := range [][2]int64{{input, inputPrice}, {output, outputPrice}, {cacheRead, cacheReadPrice}, {cacheWrite, cacheWritePrice}} { + total.Add(total, new(big.Int).Mul(big.NewInt(pair[0]), big.NewInt(pair[1]))) + } + if total.Sign() == 0 { + return 0, nil + } + total.Add(total, big.NewInt(microsPerUnit-1)) + total.Div(total, big.NewInt(microsPerUnit)) + if !total.IsInt64() { + return 0, errors.New("calculated charge exceeds supported range") + } + return total.Int64(), nil +} + +func collectibleCharge(actualCost, balance, held, reservation int64) (int64, error) { + if actualCost < 0 || balance < 0 || held < 0 || reservation < 0 || held > balance || reservation > held { + return 0, errors.New("wallet reservation invariant violated") + } + spendable := balance - (held - reservation) + if actualCost > spendable { + return spendable, nil + } + return actualCost, nil +} + +func randomHex(bytes int) string { + buffer := make([]byte, bytes) + if _, err := rand.Read(buffer); err != nil { + return fmt.Sprintf("%08x", time.Now().UnixNano())[:bytes*2] + } + return hex.EncodeToString(buffer) +} + +func randomLetters(length int) string { + const letters = "abcdefghijklmnopqrstuvwxyz" + buffer := make([]byte, length) + if _, err := rand.Read(buffer); err != nil { + return strings.Repeat("a", length) + } + for i := range buffer { + buffer[i] = letters[int(buffer[i])%len(letters)] + } + return string(buffer) +} + +func normalizeDescription(value string) string { + value = strings.TrimSpace(value) + if len(value) > 240 { + value = value[:240] + } + return value +} diff --git a/internal/billing/service_test.go b/internal/billing/service_test.go new file mode 100644 index 0000000..17f8143 --- /dev/null +++ b/internal/billing/service_test.go @@ -0,0 +1,190 @@ +package billing + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" + + "aigw/internal/controlplane" + "aigw/internal/domain" + + "github.com/stripe/stripe-go/v86" + "github.com/stripe/stripe-go/v86/webhook" +) + +func TestUsageCostUsesFixedPointAndRoundsOnce(t *testing.T) { + usage := domain.Usage{InputTokens: 3, OutputTokens: 2, CacheReadInputTokens: 5} + cost, err := usageCost(usage, 150_000, 600_000, 30_000, 0) + if err != nil { + t.Fatal(err) + } + // (3*150000 + 2*600000 + 5*30000) / 1e6 = 1.8 micros. + if cost != 2 { + t.Fatalf("cost = %d, want 2", cost) + } +} + +func TestReservationUsesExplicitOutputLimit(t *testing.T) { + model := domain.Model{InputPriceMicrosPerMillion: 1_000_000, OutputPriceMicrosPerMillion: 1_000_000} + body := []byte(`{"max_tokens":8192}`) + cost, err := reservationCost(model, body, 4096) + if err != nil { + t.Fatal(err) + } + want := int64(len(body)) + 8192 + if cost != want { + t.Fatalf("reservation = %d, want %d", cost, want) + } +} + +func TestMinorToMicrosSupportsCurrencyExponents(t *testing.T) { + tests := []struct { + currency string + minor int64 + want int64 + }{{"usd", 123, 1_230_000}, {"jpy", 123, 123_000_000}, {"bhd", 123, 123_000}} + for _, test := range tests { + got, err := minorToMicros(test.currency, test.minor) + if err != nil { + t.Fatalf("%s: %v", test.currency, err) + } + if got != test.want { + t.Fatalf("%s: got %d, want %d", test.currency, got, test.want) + } + } +} + +func TestCollectibleChargePreservesOtherReservations(t *testing.T) { + charged, err := collectibleCharge(100, 100, 80, 40) + if err != nil { + t.Fatal(err) + } + if charged != 60 { + t.Fatalf("charged = %d, want 60", charged) + } + if remainingBalance, remainingHeld := int64(100)-charged, int64(80)-40; remainingBalance < remainingHeld { + t.Fatalf("remaining balance %d does not cover held %d", remainingBalance, remainingHeld) + } +} + +func TestIntegrationIdentifierSuffixUsesLetters(t *testing.T) { + value := randomLetters(8) + if len(value) != 8 { + t.Fatalf("length = %d", len(value)) + } + for _, char := range value { + if char < 'a' || char > 'z' { + t.Fatalf("non-letter suffix %q", value) + } + } +} + +func TestWebhookRejectsInvalidSignatureBeforeProcessing(t *testing.T) { + service := &Service{stripeWebhookSecret: "whsec_test"} + request := httptest.NewRequest(http.MethodPost, "/billing/stripe/webhook", strings.NewReader(`{"id":"evt_fake"}`)) + request.Header.Set("Stripe-Signature", "invalid") + response := httptest.NewRecorder() + service.WebhookHandler().ServeHTTP(response, request) + if response.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400", response.Code) + } +} + +func TestStripeWebhookCreditsPaidOrderExactlyOncePostgres(t *testing.T) { + databaseURL := os.Getenv("AIGW_TEST_DATABASE_URL") + if databaseURL == "" { + t.Skip("AIGW_TEST_DATABASE_URL is not set") + } + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + if err := controlplane.MigrateDatabase(ctx, databaseURL); err != nil { + t.Fatal(err) + } + service, err := New(ctx, Options{ + DatabaseURL: databaseURL, Currency: "usd", MinTopUpMinor: 500, MaxTopUpMinor: 1_000_000, + StripeEnabled: true, StripeAPIKey: "rk_test_placeholder", StripeWebhookSecret: "whsec_integration_test", + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(service.Close) + + var tenantID string + slug := fmt.Sprintf("stripe-%d", time.Now().UnixNano()) + if err := service.db.QueryRow(ctx, `INSERT INTO tenants (slug, name) VALUES ($1,'Stripe integration') RETURNING id::text`, slug).Scan(&tenantID); err != nil { + t.Fatal(err) + } + eventID := fmt.Sprintf("evt_aigw_%d", time.Now().UnixNano()) + t.Cleanup(func() { + cleanupCtx := context.Background() + for _, statement := range []struct { + query string + arg string + }{ + {`DELETE FROM stripe_webhook_events WHERE event_id=$1`, eventID}, + {`DELETE FROM billing_ledger WHERE tenant_id=$1`, tenantID}, + {`DELETE FROM tenant_wallets WHERE tenant_id=$1`, tenantID}, + {`DELETE FROM topup_orders WHERE tenant_id=$1`, tenantID}, + {`DELETE FROM tenants WHERE id=$1`, tenantID}, + } { + if _, cleanupErr := service.db.Exec(cleanupCtx, statement.query, statement.arg); cleanupErr != nil { + t.Errorf("cleanup Stripe integration data: %v", cleanupErr) + } + } + }) + const amountMinor int64 = 2500 + orderID, amountMicros, err := service.createTopUpOrder(ctx, CheckoutInput{TenantID: tenantID, AmountMinor: amountMinor}) + if err != nil { + t.Fatal(err) + } + sessionID := fmt.Sprintf("cs_test_aigw_%d", time.Now().UnixNano()) + payload, err := json.Marshal(map[string]any{ + "id": eventID, "object": "event", "api_version": stripe.APIVersion, + "type": string(stripe.EventTypeCheckoutSessionCompleted), + "data": map[string]any{"object": map[string]any{ + "id": sessionID, "object": "checkout.session", "client_reference_id": orderID, + "amount_total": amountMinor, "currency": "usd", "payment_status": "paid", + }}, + }) + if err != nil { + t.Fatal(err) + } + signed := webhook.GenerateTestSignedPayload(&webhook.UnsignedPayload{Payload: payload, Secret: service.stripeWebhookSecret}) + for delivery := 0; delivery < 2; delivery++ { + request := httptest.NewRequest(http.MethodPost, "/billing/stripe/webhook", strings.NewReader(string(payload))) + request.Header.Set("Stripe-Signature", signed.Header) + response := httptest.NewRecorder() + service.WebhookHandler().ServeHTTP(response, request) + if response.Code != http.StatusOK { + t.Fatalf("delivery %d status = %d, body = %s", delivery+1, response.Code, response.Body.String()) + } + } + + var balance int64 + if err := service.db.QueryRow(ctx, `SELECT balance_micros FROM tenant_wallets WHERE tenant_id=$1`, tenantID).Scan(&balance); err != nil { + t.Fatal(err) + } + if balance != amountMicros { + t.Fatalf("balance = %d, want %d", balance, amountMicros) + } + var ledgerCount, webhookCount int + var orderStatus string + if err := service.db.QueryRow(ctx, `SELECT count(*) FROM billing_ledger WHERE source_type='stripe_checkout' AND source_id=$1`, sessionID).Scan(&ledgerCount); err != nil { + t.Fatal(err) + } + if err := service.db.QueryRow(ctx, `SELECT count(*) FROM stripe_webhook_events WHERE event_id=$1`, eventID).Scan(&webhookCount); err != nil { + t.Fatal(err) + } + if err := service.db.QueryRow(ctx, `SELECT status FROM topup_orders WHERE id=$1`, orderID).Scan(&orderStatus); err != nil { + t.Fatal(err) + } + if ledgerCount != 1 || webhookCount != 1 || orderStatus != "paid" { + t.Fatalf("ledger=%d webhook=%d order=%s, want 1/1/paid", ledgerCount, webhookCount, orderStatus) + } +} diff --git a/internal/billing/stripe.go b/internal/billing/stripe.go new file mode 100644 index 0000000..b06eb6a --- /dev/null +++ b/internal/billing/stripe.go @@ -0,0 +1,203 @@ +package billing + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strings" + + "github.com/jackc/pgx/v5" + "github.com/stripe/stripe-go/v86" + "github.com/stripe/stripe-go/v86/webhook" +) + +const maxWebhookBodyBytes = 1 << 20 + +type stripeCheckoutCreator func(context.Context, *stripe.CheckoutSessionCreateParams) (*stripe.CheckoutSession, error) + +func newStripeCheckoutCreator(apiKey string) stripeCheckoutCreator { + client := stripe.NewClient(apiKey) + return client.V1CheckoutSessions.Create +} + +func (s *Service) CreateCheckout(ctx context.Context, input CheckoutInput) (CheckoutResult, error) { + orderID, _, err := s.createTopUpOrder(ctx, input) + if err != nil { + return CheckoutResult{}, err + } + params := &stripe.CheckoutSessionCreateParams{ + Mode: stripe.String("payment"), + ClientReferenceID: stripe.String(orderID), + IntegrationIdentifier: stripe.String(s.integrationIdentifier), + SuccessURL: stripe.String(s.stripeSuccessURL), + CancelURL: stripe.String(s.stripeCancelURL), + Metadata: map[string]string{ + "aigw_topup_order_id": orderID, + "aigw_tenant_id": strings.TrimSpace(input.TenantID), + }, + LineItems: []*stripe.CheckoutSessionCreateLineItemParams{{ + Quantity: stripe.Int64(1), + PriceData: &stripe.CheckoutSessionCreateLineItemPriceDataParams{ + Currency: stripe.String(s.currency), + UnitAmount: stripe.Int64(input.AmountMinor), + ProductData: &stripe.CheckoutSessionCreateLineItemPriceDataProductDataParams{ + Name: stripe.String("AIGW prepaid balance"), + Description: stripe.String("Prepaid API usage credit"), + }, + }, + }}, + } + params.SetIdempotencyKey("aigw_topup_" + orderID) + session, err := s.createStripeCheckout(ctx, params) + if err != nil { + _, _ = s.db.Exec(ctx, `UPDATE topup_orders SET status = 'failed' WHERE id = $1 AND status = 'pending'`, orderID) + return CheckoutResult{}, fmt.Errorf("create Stripe Checkout Session: %w", err) + } + if session.ID == "" || session.URL == "" { + return CheckoutResult{}, errors.New("Stripe returned an incomplete Checkout Session") + } + if _, err := s.db.Exec(ctx, ` + UPDATE topup_orders SET stripe_session_id = $2, checkout_url = $3 + WHERE id = $1`, orderID, session.ID, session.URL); err != nil { + return CheckoutResult{}, fmt.Errorf("persist Stripe Checkout Session: %w", err) + } + return CheckoutResult{OrderID: orderID, SessionID: session.ID, URL: session.URL}, nil +} + +func (s *Service) WebhookHandler() http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + w.Header().Set("Allow", http.MethodPost) + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + body, err := io.ReadAll(io.LimitReader(r.Body, maxWebhookBodyBytes+1)) + if err != nil || len(body) > maxWebhookBodyBytes { + http.Error(w, "invalid webhook body", http.StatusBadRequest) + return + } + event, err := webhook.ConstructEvent(body, r.Header.Get("Stripe-Signature"), s.stripeWebhookSecret) + if err != nil { + http.Error(w, "invalid webhook signature", http.StatusBadRequest) + return + } + if err := s.processStripeEvent(r.Context(), event); err != nil { + if errors.Is(err, ErrInvalidAmount) || isNotFound(err) { + http.Error(w, "invalid checkout event", http.StatusBadRequest) + return + } + http.Error(w, "webhook processing failed", http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"received":true}`+"\n") + }) +} + +func (s *Service) processStripeEvent(ctx context.Context, event stripe.Event) error { + typeName := string(event.Type) + switch event.Type { + case stripe.EventTypeCheckoutSessionCompleted, + stripe.EventTypeCheckoutSessionAsyncPaymentSucceeded, + stripe.EventTypeCheckoutSessionAsyncPaymentFailed, + stripe.EventTypeCheckoutSessionExpired: + default: + return nil + } + if event.Data == nil { + return ErrInvalidAmount + } + var session stripe.CheckoutSession + if err := json.Unmarshal(event.Data.Raw, &session); err != nil { + return ErrInvalidAmount + } + if event.ID == "" || session.ID == "" || session.ClientReferenceID == "" { + return ErrInvalidAmount + } + tx, err := s.db.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.ReadCommitted}) + if err != nil { + return err + } + defer tx.Rollback(ctx) + tag, err := tx.Exec(ctx, ` + INSERT INTO stripe_webhook_events (event_id, event_type) VALUES ($1,$2) + ON CONFLICT (event_id) DO NOTHING`, event.ID, typeName) + if err != nil { + return fmt.Errorf("record Stripe event: %w", err) + } + if tag.RowsAffected() == 0 { + return tx.Commit(ctx) + } + + var tenantID, currency, status string + var amountMinor, amountMicros int64 + var storedSessionID *string + err = tx.QueryRow(ctx, ` + SELECT tenant_id::text, amount_minor, amount_micros, currency, status, stripe_session_id + FROM topup_orders WHERE id = $1 FOR UPDATE`, session.ClientReferenceID, + ).Scan(&tenantID, &amountMinor, &amountMicros, ¤cy, &status, &storedSessionID) + if err != nil { + return err + } + if (storedSessionID != nil && *storedSessionID != session.ID) || amountMinor != session.AmountTotal || currency != string(session.Currency) { + return ErrInvalidAmount + } + if event.Type == stripe.EventTypeCheckoutSessionAsyncPaymentFailed || event.Type == stripe.EventTypeCheckoutSessionExpired { + orderStatus := "failed" + if event.Type == stripe.EventTypeCheckoutSessionExpired { + orderStatus = "expired" + } + if _, err := tx.Exec(ctx, `UPDATE topup_orders SET status = $2, stripe_session_id = COALESCE(stripe_session_id, $3) WHERE id = $1 AND status = 'pending'`, session.ClientReferenceID, orderStatus, session.ID); err != nil { + return err + } + _, err = tx.Exec(ctx, `UPDATE stripe_webhook_events SET processed_at = now() WHERE event_id = $1`, event.ID) + if err != nil { + return err + } + return tx.Commit(ctx) + } + if session.PaymentStatus != stripe.CheckoutSessionPaymentStatusPaid { + _, err = tx.Exec(ctx, `UPDATE stripe_webhook_events SET processed_at = now() WHERE event_id = $1`, event.ID) + if err != nil { + return err + } + return tx.Commit(ctx) + } + if status != "paid" { + if _, err := tx.Exec(ctx, ` + INSERT INTO tenant_wallets (tenant_id, currency) VALUES ($1,$2) + ON CONFLICT (tenant_id) DO NOTHING`, tenantID, currency); err != nil { + return err + } + var walletCurrency string + var balance int64 + if err := tx.QueryRow(ctx, `SELECT currency, balance_micros FROM tenant_wallets WHERE tenant_id = $1 FOR UPDATE`, tenantID).Scan(&walletCurrency, &balance); err != nil { + return err + } + if walletCurrency != currency || amountMicros <= 0 || balance > int64(^uint64(0)>>1)-amountMicros { + return ErrInvalidAmount + } + newBalance := balance + amountMicros + if _, err := tx.Exec(ctx, `UPDATE tenant_wallets SET balance_micros = $2, updated_at = now() WHERE tenant_id = $1`, tenantID, newBalance); err != nil { + return err + } + if _, err := tx.Exec(ctx, ` + INSERT INTO billing_ledger (tenant_id, currency, amount_micros, balance_after_micros, kind, source_type, source_id, description) + VALUES ($1,$2,$3,$4,'topup','stripe_checkout',$5,'Stripe balance top-up') + ON CONFLICT (source_type, source_id) DO NOTHING`, tenantID, currency, amountMicros, newBalance, session.ID); err != nil { + return err + } + if _, err := tx.Exec(ctx, ` + UPDATE topup_orders SET status = 'paid', stripe_session_id = COALESCE(stripe_session_id, $2), paid_at = now() + WHERE id = $1`, session.ClientReferenceID, session.ID); err != nil { + return err + } + } + if _, err := tx.Exec(ctx, `UPDATE stripe_webhook_events SET processed_at = now() WHERE event_id = $1`, event.ID); err != nil { + return err + } + return tx.Commit(ctx) +} diff --git a/internal/billing/types.go b/internal/billing/types.go new file mode 100644 index 0000000..24694cc --- /dev/null +++ b/internal/billing/types.go @@ -0,0 +1,83 @@ +package billing + +import ( + "context" + "errors" + "time" + + "aigw/internal/domain" +) + +var ( + ErrInsufficientBalance = errors.New("insufficient balance") + ErrStripeDisabled = errors.New("Stripe top-ups are disabled") + ErrInvalidAmount = errors.New("invalid amount") + ErrQuotaExceeded = errors.New("monthly spend quota exceeded") +) + +type Meter interface { + Authorize(context.Context, Authorization) error + Settle(context.Context, domain.UsageEvent) error +} + +type Authorization struct { + RequestID string + Principal domain.Principal + Model domain.Model + Body []byte + Policy domain.LimitPolicy +} + +type Options struct { + DatabaseURL string + Currency string + DefaultMaxOutputTokens int64 + MinTopUpMinor int64 + MaxTopUpMinor int64 + StripeEnabled bool + StripeAPIKey string + StripeWebhookSecret string + StripeSuccessURL string + StripeCancelURL string +} + +type Account struct { + TenantID string `json:"tenant_id"` + TenantName string `json:"tenant_name"` + Currency string `json:"currency"` + BalanceMicros int64 `json:"balance_micros"` + ReservedMicros int64 `json:"reserved_micros"` + AvailableMicros int64 `json:"available_micros"` + UpdatedAt time.Time `json:"updated_at"` +} + +type LedgerEntry struct { + ID string `json:"id"` + TenantID string `json:"tenant_id"` + ProjectID string `json:"project_id,omitempty"` + Currency string `json:"currency"` + AmountMicros int64 `json:"amount_micros"` + BalanceAfterMicros int64 `json:"balance_after_micros"` + Kind string `json:"kind"` + SourceType string `json:"source_type"` + SourceID string `json:"source_id"` + Description string `json:"description"` + CreatedAt time.Time `json:"created_at"` +} + +type AdjustmentInput struct { + TenantID string `json:"tenant_id"` + AmountMicros int64 `json:"amount_micros"` + Description string `json:"description"` +} + +type CheckoutInput struct { + TenantID string `json:"tenant_id"` + AmountMinor int64 `json:"amount_minor"` +} + +type CheckoutResult struct { + OrderID string `json:"order_id"` + SessionID string `json:"session_id"` + URL string `json:"url"` +} diff --git a/internal/catalog/catalog.go b/internal/catalog/catalog.go index 76e33b1..7966d9a 100644 --- a/internal/catalog/catalog.go +++ b/internal/catalog/catalog.go @@ -32,7 +32,14 @@ func New(cfg config.Config) *Catalog { models := make([]domain.Model, 0, len(cfg.Models)) for _, modelCfg := range cfg.Models { - model := domain.Model{ID: modelCfg.ID, OwnedBy: modelCfg.OwnedBy} + model := domain.Model{ + ID: modelCfg.ID, + OwnedBy: modelCfg.OwnedBy, + InputPriceMicrosPerMillion: modelCfg.InputPriceMicrosPerMillion, + OutputPriceMicrosPerMillion: modelCfg.OutputPriceMicrosPerMillion, + CacheReadPriceMicrosPerMillion: modelCfg.CacheReadPriceMicrosPerMillion, + CacheWritePriceMicrosPerMillion: modelCfg.CacheWritePriceMicrosPerMillion, + } for _, route := range modelCfg.Routes { model.Routes = append(model.Routes, domain.Route{ Provider: providers[route.Provider], diff --git a/internal/config/config.go b/internal/config/config.go index 5e1e518..8c21c6c 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -21,6 +21,7 @@ type Config struct { UpstreamHTTP UpstreamHTTPConfig `json:"upstream_http"` Providers []ProviderConfig `json:"providers"` Models []ModelConfig `json:"models"` + Billing BillingConfig `json:"billing"` Observability ObservabilityConfig `json:"observability"` } @@ -74,9 +75,13 @@ type ProviderConfig struct { } type ModelConfig struct { - ID string `json:"id"` - OwnedBy string `json:"owned_by"` - Routes []RouteConfig `json:"routes"` + ID string `json:"id"` + OwnedBy string `json:"owned_by"` + InputPriceMicrosPerMillion int64 `json:"input_price_micros_per_million"` + OutputPriceMicrosPerMillion int64 `json:"output_price_micros_per_million"` + CacheReadPriceMicrosPerMillion int64 `json:"cache_read_price_micros_per_million"` + CacheWritePriceMicrosPerMillion int64 `json:"cache_write_price_micros_per_million"` + Routes []RouteConfig `json:"routes"` } type RouteConfig struct { @@ -91,6 +96,25 @@ type ObservabilityConfig struct { ExposeMetrics bool `json:"expose_metrics"` } +type BillingConfig struct { + Enabled bool `json:"enabled"` + Currency string `json:"currency"` + DefaultMaxOutputTokens int64 `json:"default_max_output_tokens"` + MinTopUpMinor int64 `json:"min_top_up_minor"` + MaxTopUpMinor int64 `json:"max_top_up_minor"` + Stripe StripeConfig `json:"stripe"` +} + +type StripeConfig struct { + Enabled bool `json:"enabled"` + APIKeyEnv string `json:"api_key_env"` + WebhookSecretEnv string `json:"webhook_secret_env"` + SuccessURL string `json:"success_url"` + CancelURL string `json:"cancel_url"` + APIKey string `json:"-"` + WebhookSecret string `json:"-"` +} + func Load(path string) (Config, error) { f, err := os.Open(path) if err != nil { @@ -179,6 +203,24 @@ func applyDefaults(cfg *Config) { if cfg.Observability.UsageBuffer == 0 { cfg.Observability.UsageBuffer = 8192 } + if cfg.Billing.Currency == "" { + cfg.Billing.Currency = "usd" + } + if cfg.Billing.DefaultMaxOutputTokens == 0 { + cfg.Billing.DefaultMaxOutputTokens = 4096 + } + if cfg.Billing.MaxTopUpMinor == 0 { + cfg.Billing.MaxTopUpMinor = 1000000 + } + if cfg.Billing.MinTopUpMinor == 0 { + cfg.Billing.MinTopUpMinor = 500 + } + if cfg.Billing.Stripe.APIKeyEnv == "" { + cfg.Billing.Stripe.APIKeyEnv = "AIGW_STRIPE_API_KEY" + } + if cfg.Billing.Stripe.WebhookSecretEnv == "" { + cfg.Billing.Stripe.WebhookSecretEnv = "AIGW_STRIPE_WEBHOOK_SECRET" + } for i := range cfg.Models { for j := range cfg.Models[i].Routes { if cfg.Models[i].Routes[j].Weight == 0 { @@ -197,6 +239,10 @@ func resolveSecrets(cfg *Config) error { if cfg.Admin.Enabled { cfg.Admin.Token = os.Getenv(cfg.Admin.TokenEnv) } + if cfg.Billing.Enabled && cfg.Billing.Stripe.Enabled { + cfg.Billing.Stripe.APIKey = os.Getenv(cfg.Billing.Stripe.APIKeyEnv) + cfg.Billing.Stripe.WebhookSecret = os.Getenv(cfg.Billing.Stripe.WebhookSecretEnv) + } for i := range cfg.Providers { provider := &cfg.Providers[i] if provider.APIKeyEnv == "" { @@ -240,6 +286,40 @@ func Validate(cfg Config) error { return errors.New("admin.base_path must start with / and cannot be /") } } + if cfg.Billing.Enabled { + if !cfg.ControlPlane.Enabled { + return errors.New("billing requires control_plane.enabled") + } + if cfg.Auth.AllowAnonymous { + return errors.New("billing cannot be enabled with auth.allow_anonymous") + } + if len(cfg.Billing.Currency) != 3 || strings.ToLower(cfg.Billing.Currency) != cfg.Billing.Currency { + return errors.New("billing.currency must be a lowercase ISO 4217 code") + } + if cfg.Billing.DefaultMaxOutputTokens < 1 { + return errors.New("billing.default_max_output_tokens must be positive") + } + if cfg.Billing.MinTopUpMinor < 1 || cfg.Billing.MaxTopUpMinor < cfg.Billing.MinTopUpMinor { + return errors.New("billing top-up bounds are invalid") + } + if cfg.Billing.Stripe.Enabled { + if cfg.Billing.Stripe.APIKey == "" { + return fmt.Errorf("billing.stripe: environment variable %s is empty", cfg.Billing.Stripe.APIKeyEnv) + } + if cfg.Billing.Stripe.WebhookSecret == "" { + return fmt.Errorf("billing.stripe: environment variable %s is empty", cfg.Billing.Stripe.WebhookSecretEnv) + } + if strings.TrimSpace(cfg.Billing.Stripe.SuccessURL) == "" || strings.TrimSpace(cfg.Billing.Stripe.CancelURL) == "" { + return errors.New("billing.stripe.success_url and cancel_url are required") + } + for name, value := range map[string]string{"success_url": cfg.Billing.Stripe.SuccessURL, "cancel_url": cfg.Billing.Stripe.CancelURL} { + parsed, err := url.Parse(value) + if err != nil || parsed.Host == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") { + return fmt.Errorf("billing.stripe.%s must be an absolute http(s) URL", name) + } + } + } + } providers := make(map[string]ProviderConfig, len(cfg.Providers)) for _, provider := range cfg.Providers { @@ -274,6 +354,9 @@ func Validate(cfg Config) error { return fmt.Errorf("duplicate model id %q", model.ID) } models[model.ID] = struct{}{} + if model.InputPriceMicrosPerMillion < 0 || model.OutputPriceMicrosPerMillion < 0 || model.CacheReadPriceMicrosPerMillion < 0 || model.CacheWritePriceMicrosPerMillion < 0 { + return fmt.Errorf("model %q: prices cannot be negative", model.ID) + } if len(model.Routes) == 0 { return fmt.Errorf("model %q: at least one route is required", model.ID) } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index e681508..2331b0f 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -85,6 +85,31 @@ func TestLoadRejectsAdminWithoutControlPlane(t *testing.T) { } } +func TestLoadResolvesStripeSecrets(t *testing.T) { + t.Setenv("AIGW_DATABASE_URL", "postgres://aigw:aigw@postgres/aigw") + t.Setenv("AIGW_CREDENTIAL_KEY", "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=") + t.Setenv("TEST_STRIPE_KEY", "rk_test_example") + t.Setenv("TEST_STRIPE_WEBHOOK", "whsec_example") + path := writeConfig(t, `{ + "control_plane":{"enabled":true}, + "billing":{"enabled":true,"stripe":{"enabled":true,"api_key_env":"TEST_STRIPE_KEY","webhook_secret_env":"TEST_STRIPE_WEBHOOK","success_url":"http://localhost/admin/?topup=success","cancel_url":"http://localhost/admin/?topup=cancel"}} +}`) + cfg, err := Load(path) + if err != nil { + t.Fatal(err) + } + if cfg.Billing.Stripe.APIKey != "rk_test_example" || cfg.Billing.Stripe.WebhookSecret != "whsec_example" { + t.Fatal("Stripe secrets were not resolved") + } +} + +func TestLoadRejectsBillingWithoutControlPlane(t *testing.T) { + path := writeConfig(t, `{"billing":{"enabled":true}}`) + if _, err := Load(path); err == nil { + t.Fatal("expected billing without control plane to be rejected") + } +} + func writeConfig(t *testing.T, content string) string { t.Helper() path := filepath.Join(t.TempDir(), "config.json") diff --git a/internal/controlplane/access.go b/internal/controlplane/access.go new file mode 100644 index 0000000..ec2d177 --- /dev/null +++ b/internal/controlplane/access.go @@ -0,0 +1,156 @@ +package controlplane + +import ( + "context" + "crypto/rand" + "crypto/sha256" + "encoding/base64" + "errors" + "fmt" + "net/mail" + "strings" + + "github.com/jackc/pgx/v5" +) + +var ErrConsoleUnauthorized = errors.New("invalid console token") + +const ( + RolePlatformAdmin = "platform_admin" + RolePlatformViewer = "platform_viewer" + RoleTenantAdmin = "tenant_admin" + RoleTenantBilling = "tenant_billing" + RoleTenantDeveloper = "tenant_developer" + RoleTenantViewer = "tenant_viewer" +) + +func (a ConsoleActor) IsPlatform() bool { return strings.HasPrefix(a.Role, "platform_") } + +func (a ConsoleActor) Can(permission string) bool { + if a.Role == RolePlatformAdmin { + return true + } + read := strings.HasSuffix(permission, ".read") + switch a.Role { + case RolePlatformViewer: + return read + case RoleTenantAdmin: + switch permission { + case "overview.read", "tenants.read", "projects.read", "projects.write", "keys.read", "keys.write", + "billing.read", "billing.topup", "usage.read", "audit.read", "limits.read", "users.read", "users.write": + return true + } + return false + case RoleTenantBilling: + return permission == "overview.read" || permission == "billing.read" || permission == "billing.topup" || permission == "usage.read" || permission == "audit.read" + case RoleTenantDeveloper: + return permission == "overview.read" || permission == "tenants.read" || permission == "projects.read" || permission == "keys.read" || permission == "keys.write" || permission == "usage.read" || permission == "limits.read" + case RoleTenantViewer: + return permission == "overview.read" || permission == "tenants.read" || permission == "projects.read" || permission == "keys.read" || permission == "billing.read" || permission == "usage.read" || permission == "limits.read" || permission == "audit.read" + default: + return false + } +} + +func (a ConsoleActor) Permissions() []string { + all := []string{"overview.read", "tenants.read", "tenants.write", "projects.read", "projects.write", "keys.read", "keys.write", "platform.read", "platform.write", "billing.read", "billing.topup", "billing.adjust", "usage.read", "limits.read", "limits.write", "users.read", "users.write", "audit.read"} + result := make([]string, 0, len(all)) + for _, permission := range all { + if a.Can(permission) { + result = append(result, permission) + } + } + return result +} + +func (s *Store) AuthenticateConsoleToken(ctx context.Context, raw string) (ConsoleActor, error) { + hash := sha256.Sum256([]byte(raw)) + var actor ConsoleActor + err := s.db.QueryRow(ctx, ` + UPDATE console_users SET last_used_at = now() + WHERE token_hash = $1 AND status = 'active' + RETURNING id::text, COALESCE(tenant_id::text, ''), email, display_name, role`, hash[:], + ).Scan(&actor.ID, &actor.TenantID, &actor.Email, &actor.DisplayName, &actor.Role) + if errors.Is(err, pgx.ErrNoRows) { + return ConsoleActor{}, ErrConsoleUnauthorized + } + if err != nil { + return ConsoleActor{}, fmt.Errorf("authenticate console token: %w", err) + } + return actor, nil +} + +func (s *Store) ListConsoleUsers(ctx context.Context, tenantID string) ([]ConsoleUser, error) { + query := `SELECT id::text, COALESCE(tenant_id::text, ''), email, display_name, role, token_prefix, status, last_used_at, created_at FROM console_users` + args := []any{} + if tenantID != "" { + query += ` WHERE tenant_id = $1` + args = append(args, tenantID) + } + query += ` ORDER BY created_at DESC` + rows, err := s.db.Query(ctx, query, args...) + if err != nil { + return nil, fmt.Errorf("query console users: %w", err) + } + defer rows.Close() + result := make([]ConsoleUser, 0) + for rows.Next() { + var item ConsoleUser + if err := rows.Scan(&item.ID, &item.TenantID, &item.Email, &item.DisplayName, &item.Role, &item.TokenPrefix, &item.Status, &item.LastUsedAt, &item.CreatedAt); err != nil { + return nil, fmt.Errorf("scan console user: %w", err) + } + result = append(result, item) + } + return result, rows.Err() +} + +func (s *Store) CreateConsoleUser(ctx context.Context, input CreateConsoleUserInput) (CreatedConsoleUser, error) { + input.Email = strings.ToLower(strings.TrimSpace(input.Email)) + input.DisplayName = strings.TrimSpace(input.DisplayName) + input.TenantID = strings.TrimSpace(input.TenantID) + address, addressErr := mail.ParseAddress(input.Email) + if addressErr != nil || address.Address != input.Email || input.DisplayName == "" { + return CreatedConsoleUser{}, errors.New("console user requires a valid email and display_name") + } + platform := input.Role == RolePlatformAdmin || input.Role == RolePlatformViewer + tenant := input.Role == RoleTenantAdmin || input.Role == RoleTenantBilling || input.Role == RoleTenantDeveloper || input.Role == RoleTenantViewer + if (!platform && !tenant) || (platform && input.TenantID != "") || (tenant && input.TenantID == "") { + return CreatedConsoleUser{}, errors.New("console user role and tenant_id are inconsistent") + } + random := make([]byte, 32) + if _, err := rand.Read(random); err != nil { + return CreatedConsoleUser{}, fmt.Errorf("generate console token: %w", err) + } + raw := "cu-aigw-" + base64.RawURLEncoding.EncodeToString(random) + hash := sha256.Sum256([]byte(raw)) + prefix := raw[:min(18, len(raw))] + "..." + var result CreatedConsoleUser + err := s.db.QueryRow(ctx, ` + INSERT INTO console_users (tenant_id, email, display_name, role, token_prefix, token_hash) + VALUES (NULLIF($1,'')::uuid,$2,$3,$4,$5,$6) + RETURNING id::text, COALESCE(tenant_id::text, ''), email, display_name, role, token_prefix, status, last_used_at, created_at`, + input.TenantID, input.Email, input.DisplayName, input.Role, prefix, hash[:], + ).Scan(&result.ID, &result.TenantID, &result.Email, &result.DisplayName, &result.Role, &result.TokenPrefix, &result.Status, &result.LastUsedAt, &result.CreatedAt) + if err != nil { + return CreatedConsoleUser{}, fmt.Errorf("create console user: %w", err) + } + result.Token = raw + return result, nil +} + +func (s *Store) RevokeConsoleUser(ctx context.Context, id, tenantID string) error { + query := `UPDATE console_users SET status='revoked', revoked_at=now() WHERE id=$1 AND status='active'` + args := []any{id} + if tenantID != "" { + query += ` AND tenant_id=$2` + args = append(args, tenantID) + } + result, err := s.db.Exec(ctx, query, args...) + if err != nil { + return err + } + if result.RowsAffected() == 0 { + return ErrNotFound + } + return nil +} diff --git a/internal/controlplane/access_test.go b/internal/controlplane/access_test.go new file mode 100644 index 0000000..96e3869 --- /dev/null +++ b/internal/controlplane/access_test.go @@ -0,0 +1,39 @@ +package controlplane + +import "testing" + +func TestConsoleRolePermissions(t *testing.T) { + tests := []struct { + role, permission string + want bool + }{ + {RolePlatformAdmin, "platform.write", true}, + {RolePlatformViewer, "platform.read", true}, + {RolePlatformViewer, "platform.write", false}, + {RoleTenantAdmin, "keys.write", true}, + {RoleTenantAdmin, "limits.write", false}, + {RoleTenantBilling, "billing.topup", true}, + {RoleTenantBilling, "keys.read", false}, + {RoleTenantDeveloper, "keys.write", true}, + {RoleTenantDeveloper, "billing.read", false}, + {RoleTenantViewer, "usage.read", true}, + {RoleTenantViewer, "users.read", false}, + } + for _, test := range tests { + t.Run(test.role+"/"+test.permission, func(t *testing.T) { + if got := (ConsoleActor{Role: test.role}).Can(test.permission); got != test.want { + t.Fatalf("Can(%q) = %v, want %v", test.permission, got, test.want) + } + }) + } +} + +func TestTenantRoleNeverGetsPlatformPermissions(t *testing.T) { + roles := []string{RoleTenantAdmin, RoleTenantBilling, RoleTenantDeveloper, RoleTenantViewer} + for _, role := range roles { + actor := ConsoleActor{Role: role} + if actor.Can("platform.read") || actor.Can("platform.write") { + t.Fatalf("role %s received platform access", role) + } + } +} diff --git a/internal/controlplane/audit.go b/internal/controlplane/audit.go new file mode 100644 index 0000000..93a4f58 --- /dev/null +++ b/internal/controlplane/audit.go @@ -0,0 +1,52 @@ +package controlplane + +import ( + "context" + "fmt" + "strings" +) + +func (s *Store) WriteAudit(ctx context.Context, input AuditInput) error { + actorType := "console_user" + if input.Actor.Bootstrap { + actorType = "bootstrap" + } + _, err := s.db.Exec(ctx, `INSERT INTO audit_logs + (actor_id, actor_type, actor_role, tenant_id, request_id, method, path, action, status_code, remote_ip, user_agent) + VALUES (NULLIF($1,'')::uuid,$2,$3,NULLIF($4,'')::uuid,$5,$6,$7,$8,$9,NULLIF($10,'')::inet,$11)`, + input.Actor.ID, actorType, input.Actor.Role, input.Actor.TenantID, input.RequestID, input.Method, + input.Path, input.Action, input.StatusCode, input.RemoteIP, input.UserAgent) + if err != nil { + return fmt.Errorf("write audit log: %w", err) + } + return nil +} + +func (s *Store) ListAudit(ctx context.Context, tenantID string, limit int) ([]AuditLog, error) { + if limit < 1 || limit > 1000 { + limit = 200 + } + query := `SELECT id, COALESCE(actor_id::text,''), actor_type, actor_role, COALESCE(tenant_id::text,''), + request_id, method, path, action, status_code, COALESCE(host(remote_ip),''), user_agent, created_at FROM audit_logs` + args := []any{} + if strings.TrimSpace(tenantID) != "" { + query += ` WHERE tenant_id=$1` + args = append(args, tenantID) + } + args = append(args, limit) + query += fmt.Sprintf(` ORDER BY created_at DESC LIMIT $%d`, len(args)) + rows, err := s.db.Query(ctx, query, args...) + if err != nil { + return nil, fmt.Errorf("query audit logs: %w", err) + } + defer rows.Close() + result := make([]AuditLog, 0) + for rows.Next() { + var item AuditLog + if err := rows.Scan(&item.ID, &item.ActorID, &item.ActorType, &item.ActorRole, &item.TenantID, &item.RequestID, &item.Method, &item.Path, &item.Action, &item.StatusCode, &item.RemoteIP, &item.UserAgent, &item.CreatedAt); err != nil { + return nil, fmt.Errorf("scan audit log: %w", err) + } + result = append(result, item) + } + return result, rows.Err() +} diff --git a/internal/controlplane/limits.go b/internal/controlplane/limits.go new file mode 100644 index 0000000..5b5825e --- /dev/null +++ b/internal/controlplane/limits.go @@ -0,0 +1,81 @@ +package controlplane + +import ( + "context" + "errors" + "fmt" + "strings" + + "github.com/jackc/pgx/v5" +) + +func (s *Store) ListProjectLimits(ctx context.Context, tenantID string) ([]ProjectLimit, error) { + query := `SELECT l.tenant_id::text, l.project_id::text, p.name, t.name, + l.requests_per_minute, l.tokens_per_minute, l.concurrent_requests, l.monthly_spend_micros, l.updated_at + FROM project_limits l JOIN projects p ON p.id=l.project_id JOIN tenants t ON t.id=l.tenant_id` + args := []any{} + if strings.TrimSpace(tenantID) != "" { + query += ` WHERE l.tenant_id=$1` + args = append(args, tenantID) + } + query += ` ORDER BY t.name, p.name` + rows, err := s.db.Query(ctx, query, args...) + if err != nil { + return nil, fmt.Errorf("query project limits: %w", err) + } + defer rows.Close() + result := make([]ProjectLimit, 0) + for rows.Next() { + var item ProjectLimit + if err := rows.Scan(&item.TenantID, &item.ProjectID, &item.ProjectName, &item.TenantName, + &item.RequestsPerMinute, &item.TokensPerMinute, &item.Concurrent, &item.MonthlySpendMicros, &item.UpdatedAt); err != nil { + return nil, fmt.Errorf("scan project limit: %w", err) + } + result = append(result, item) + } + return result, rows.Err() +} + +func (s *Store) SetProjectLimit(ctx context.Context, projectID string, input SetProjectLimitInput) (ProjectLimit, int64, error) { + if strings.TrimSpace(projectID) == "" || input.RequestsPerMinute < 0 || input.TokensPerMinute < 0 || input.ConcurrentRequests < 0 || input.MonthlySpendMicros < 0 { + return ProjectLimit{}, 0, errors.New("limit values cannot be negative and project_id is required") + } + tx, err := s.db.Begin(ctx) + if err != nil { + return ProjectLimit{}, 0, err + } + defer tx.Rollback(ctx) + var result ProjectLimit + err = tx.QueryRow(ctx, ` + WITH project AS ( + SELECT p.id, p.tenant_id, p.name AS project_name, t.name AS tenant_name + FROM projects p JOIN tenants t ON t.id=p.tenant_id WHERE p.id=$1 + ), updated AS ( + INSERT INTO project_limits (project_id, tenant_id, requests_per_minute, tokens_per_minute, concurrent_requests, monthly_spend_micros) + SELECT p.id, p.tenant_id, $2, $3, $4, $5 FROM project p + ON CONFLICT (project_id) DO UPDATE SET requests_per_minute=EXCLUDED.requests_per_minute, + tokens_per_minute=EXCLUDED.tokens_per_minute, concurrent_requests=EXCLUDED.concurrent_requests, + monthly_spend_micros=EXCLUDED.monthly_spend_micros, updated_at=now() + RETURNING tenant_id, project_id, requests_per_minute, tokens_per_minute, concurrent_requests, monthly_spend_micros, updated_at + ) + SELECT u.tenant_id::text, u.project_id::text, p.project_name, p.tenant_name, + u.requests_per_minute, u.tokens_per_minute, u.concurrent_requests, u.monthly_spend_micros, u.updated_at + FROM updated u JOIN project p ON p.id=u.project_id`, + projectID, input.RequestsPerMinute, input.TokensPerMinute, input.ConcurrentRequests, input.MonthlySpendMicros, + ).Scan(&result.TenantID, &result.ProjectID, &result.ProjectName, &result.TenantName, + &result.RequestsPerMinute, &result.TokensPerMinute, &result.Concurrent, &result.MonthlySpendMicros, &result.UpdatedAt) + if errors.Is(err, pgx.ErrNoRows) { + return ProjectLimit{}, 0, ErrNotFound + } + if err != nil { + return ProjectLimit{}, 0, fmt.Errorf("set project limit: %w", err) + } + generation, err := bumpGeneration(ctx, tx) + if err != nil { + return ProjectLimit{}, 0, err + } + if err := tx.Commit(ctx); err != nil { + return ProjectLimit{}, 0, err + } + return result, generation, nil +} diff --git a/internal/controlplane/limits_integration_test.go b/internal/controlplane/limits_integration_test.go new file mode 100644 index 0000000..8087bdd --- /dev/null +++ b/internal/controlplane/limits_integration_test.go @@ -0,0 +1,62 @@ +package controlplane + +import ( + "context" + "fmt" + "os" + "testing" + "time" +) + +func TestSetProjectLimitPostgres(t *testing.T) { + databaseURL := os.Getenv("AIGW_TEST_DATABASE_URL") + if databaseURL == "" { + t.Skip("AIGW_TEST_DATABASE_URL is not set") + } + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + if err := MigrateDatabase(ctx, databaseURL); err != nil { + t.Fatal(err) + } + store, err := NewStore(ctx, Options{ + DatabaseURL: databaseURL, + CredentialKey: "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=", + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if err := store.Close(); err != nil { + t.Errorf("close control-plane store: %v", err) + } + }) + + suffix := time.Now().UnixNano() + tenant, _, err := store.CreateTenant(ctx, CreateTenantInput{Slug: fmt.Sprintf("limits-%d", suffix), Name: "Limits tenant"}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if _, cleanupErr := store.db.Exec(context.Background(), `DELETE FROM tenants WHERE id=$1`, tenant.ID); cleanupErr != nil { + t.Errorf("cleanup limit integration data: %v", cleanupErr) + } + }) + project, _, err := store.CreateProject(ctx, CreateProjectInput{TenantID: tenant.ID, Slug: "limits-project", Name: "Limits project"}) + if err != nil { + t.Fatal(err) + } + + got, _, err := store.SetProjectLimit(ctx, project.ID, SetProjectLimitInput{ + RequestsPerMinute: 10, TokensPerMinute: 20, ConcurrentRequests: 3, MonthlySpendMicros: 40, + }) + if err != nil { + t.Fatal(err) + } + if got.TenantID != tenant.ID || got.ProjectID != project.ID || got.TenantName != tenant.Name || got.ProjectName != project.Name { + t.Fatalf("unexpected limit identity: %+v", got) + } + if got.RequestsPerMinute != 10 || got.TokensPerMinute != 20 || got.Concurrent != 3 || got.MonthlySpendMicros != 40 { + t.Fatalf("unexpected limit values: %+v", got) + } +} diff --git a/internal/controlplane/manager.go b/internal/controlplane/manager.go index 42f5efe..cdafc36 100644 --- a/internal/controlplane/manager.go +++ b/internal/controlplane/manager.go @@ -10,6 +10,7 @@ import ( "aigw/internal/auth" "aigw/internal/catalog" + "aigw/internal/domain" ) const broadcastQueueSize = 128 @@ -22,6 +23,10 @@ type managerStore interface { RedisEnabled() bool } +type policyReplacer interface { + ReplacePolicies([]domain.LimitPolicy) +} + type Manager struct { store managerStore catalog *catalog.Catalog @@ -32,15 +37,20 @@ type Manager struct { redisConnected atomic.Bool reloadMu sync.Mutex broadcasts chan ChangeEvent + policyTarget policyReplacer } -func NewManager(store managerStore, modelCatalog *catalog.Catalog, authenticator *auth.StaticAuthenticator, logger *slog.Logger, pollInterval time.Duration) *Manager { +func NewManager(store managerStore, modelCatalog *catalog.Catalog, authenticator *auth.StaticAuthenticator, logger *slog.Logger, pollInterval time.Duration, policyTargets ...policyReplacer) *Manager { if pollInterval <= 0 { pollInterval = 30 * time.Second } + var policyTarget policyReplacer + if len(policyTargets) > 0 { + policyTarget = policyTargets[0] + } return &Manager{ store: store, catalog: modelCatalog, authenticator: authenticator, - logger: logger, pollInterval: pollInterval, broadcasts: make(chan ChangeEvent, broadcastQueueSize), + logger: logger, pollInterval: pollInterval, broadcasts: make(chan ChangeEvent, broadcastQueueSize), policyTarget: policyTarget, } } @@ -53,6 +63,9 @@ func (m *Manager) Reload(ctx context.Context) (int64, error) { } m.catalog.Replace(snapshot.Models) m.authenticator.ReplaceHashed(snapshot.APIKeys) + if m.policyTarget != nil { + m.policyTarget.ReplacePolicies(snapshot.Limits) + } m.generation.Store(snapshot.Generation) m.logger.Info("control_plane_reloaded", "generation", snapshot.Generation, "models", len(snapshot.Models), "api_keys", len(snapshot.APIKeys)) return snapshot.Generation, nil @@ -195,7 +208,10 @@ func (m *Manager) runBroadcasts(ctx context.Context) { err := m.store.PublishChange(publishContext, event) cancel() if err != nil { + m.redisConnected.Store(false) m.logger.Warn("control_plane_publish_failed", "generation", event.Generation, "resource", event.Resource, "id", event.ID, "error", err) + } else { + m.redisConnected.Store(true) } } } diff --git a/internal/controlplane/manager_test.go b/internal/controlplane/manager_test.go index 8dd4012..ee78a00 100644 --- a/internal/controlplane/manager_test.go +++ b/internal/controlplane/manager_test.go @@ -13,6 +13,7 @@ import ( "aigw/internal/auth" "aigw/internal/catalog" + "aigw/internal/domain" ) type fakeManagerStore struct { @@ -26,6 +27,12 @@ type fakeManagerStore struct { subscribe func(context.Context, int64) (<-chan ChangeMessage, func() error, error) } +type capturePolicies struct{ values []domain.LimitPolicy } + +func (c *capturePolicies) ReplacePolicies(values []domain.LimitPolicy) { + c.values = append([]domain.LimitPolicy(nil), values...) +} + type safeLogBuffer struct { mu sync.Mutex buf bytes.Buffer @@ -199,6 +206,19 @@ func TestRedisCanBeDisabled(t *testing.T) { } } +func TestReloadReplacesLimitPolicySnapshot(t *testing.T) { + store := newFakeManagerStore(4) + store.snapshot.Store(Snapshot{Generation: 4, Limits: []domain.LimitPolicy{{ProjectID: "project-1", Concurrent: 3}}}) + target := &capturePolicies{} + manager := NewManager(store, catalog.NewModels(nil), auth.NewDynamic(nil, false), slog.Default(), time.Second, target) + if _, err := manager.Reload(context.Background()); err != nil { + t.Fatal(err) + } + if len(target.values) != 1 || target.values[0].Concurrent != 3 { + t.Fatalf("unexpected policies: %+v", target.values) + } +} + func waitUntil(t *testing.T, timeout time.Duration, condition func() bool) { t.Helper() deadline := time.Now().Add(timeout) diff --git a/internal/controlplane/mutations.go b/internal/controlplane/mutations.go index a781994..3cedf70 100644 --- a/internal/controlplane/mutations.go +++ b/internal/controlplane/mutations.go @@ -176,6 +176,9 @@ func (s *Store) CreateModel(ctx context.Context, input CreateModelInput) (Model, if input.PublicID == "" || len(input.Routes) == 0 { return Model{}, 0, errors.New("model requires public_id and at least one route") } + if input.InputPriceMicrosPerMillion < 0 || input.OutputPriceMicrosPerMillion < 0 || input.CacheReadPriceMicrosPerMillion < 0 || input.CacheWritePriceMicrosPerMillion < 0 { + return Model{}, 0, errors.New("model prices cannot be negative") + } for i := range input.Routes { input.Routes[i].ProviderID = strings.TrimSpace(input.Routes[i].ProviderID) input.Routes[i].UpstreamModel = strings.TrimSpace(input.Routes[i].UpstreamModel) @@ -193,9 +196,16 @@ func (s *Store) CreateModel(ctx context.Context, input CreateModelInput) (Model, defer tx.Rollback(ctx) var result Model err = tx.QueryRow(ctx, ` - INSERT INTO models (public_id, owned_by) VALUES ($1, $2) - RETURNING id::text, public_id, owned_by, enabled, created_at`, input.PublicID, input.OwnedBy, - ).Scan(&result.ID, &result.PublicID, &result.OwnedBy, &result.Enabled, &result.CreatedAt) + INSERT INTO models (public_id, owned_by, input_price_micros_per_million, output_price_micros_per_million, + cache_read_price_micros_per_million, cache_write_price_micros_per_million) + VALUES ($1, $2, $3, $4, $5, $6) + RETURNING id::text, public_id, owned_by, input_price_micros_per_million, output_price_micros_per_million, + cache_read_price_micros_per_million, cache_write_price_micros_per_million, enabled, created_at`, + input.PublicID, input.OwnedBy, input.InputPriceMicrosPerMillion, input.OutputPriceMicrosPerMillion, + input.CacheReadPriceMicrosPerMillion, input.CacheWritePriceMicrosPerMillion, + ).Scan(&result.ID, &result.PublicID, &result.OwnedBy, &result.InputPriceMicrosPerMillion, + &result.OutputPriceMicrosPerMillion, &result.CacheReadPriceMicrosPerMillion, + &result.CacheWritePriceMicrosPerMillion, &result.Enabled, &result.CreatedAt) if err != nil { return Model{}, 0, fmt.Errorf("create model: %w", err) } diff --git a/internal/controlplane/queries.go b/internal/controlplane/queries.go index 732cb46..9b76fad 100644 --- a/internal/controlplane/queries.go +++ b/internal/controlplane/queries.go @@ -24,7 +24,18 @@ func (s *Store) Overview(ctx context.Context) (Overview, error) { } func (s *Store) ListTenants(ctx context.Context) ([]Tenant, error) { - rows, err := s.db.Query(ctx, `SELECT id::text, slug, name, status, created_at FROM tenants ORDER BY created_at DESC`) + return s.ListTenantsFor(ctx, "") +} + +func (s *Store) ListTenantsFor(ctx context.Context, tenantID string) ([]Tenant, error) { + query := `SELECT id::text, slug, name, status, created_at FROM tenants` + args := []any{} + if tenantID != "" { + query += ` WHERE id=$1` + args = append(args, tenantID) + } + query += ` ORDER BY created_at DESC` + rows, err := s.db.Query(ctx, query, args...) if err != nil { return nil, fmt.Errorf("query tenants: %w", err) } @@ -41,7 +52,18 @@ func (s *Store) ListTenants(ctx context.Context) ([]Tenant, error) { } func (s *Store) ListProjects(ctx context.Context) ([]Project, error) { - rows, err := s.db.Query(ctx, `SELECT id::text, tenant_id::text, slug, name, status, created_at FROM projects ORDER BY created_at DESC`) + return s.ListProjectsFor(ctx, "") +} + +func (s *Store) ListProjectsFor(ctx context.Context, tenantID string) ([]Project, error) { + query := `SELECT id::text, tenant_id::text, slug, name, status, created_at FROM projects` + args := []any{} + if tenantID != "" { + query += ` WHERE tenant_id=$1` + args = append(args, tenantID) + } + query += ` ORDER BY created_at DESC` + rows, err := s.db.Query(ctx, query, args...) if err != nil { return nil, fmt.Errorf("query projects: %w", err) } @@ -58,9 +80,20 @@ func (s *Store) ListProjects(ctx context.Context) ([]Project, error) { } func (s *Store) ListAPIKeys(ctx context.Context) ([]APIKey, error) { - rows, err := s.db.Query(ctx, ` + return s.ListAPIKeysFor(ctx, "") +} + +func (s *Store) ListAPIKeysFor(ctx context.Context, tenantID string) ([]APIKey, error) { + query := ` SELECT id::text, tenant_id::text, project_id::text, name, key_prefix, scopes, status, created_at - FROM api_keys ORDER BY created_at DESC`) + FROM api_keys` + args := []any{} + if tenantID != "" { + query += ` WHERE tenant_id=$1` + args = append(args, tenantID) + } + query += ` ORDER BY created_at DESC` + rows, err := s.db.Query(ctx, query, args...) if err != nil { return nil, fmt.Errorf("query API keys: %w", err) } @@ -80,6 +113,44 @@ func (s *Store) ListAPIKeys(ctx context.Context) ([]APIKey, error) { return result, rows.Err() } +func (s *Store) OverviewFor(ctx context.Context, tenantID string) (Overview, error) { + if tenantID == "" { + return s.Overview(ctx) + } + var result Overview + err := s.db.QueryRow(ctx, `SELECT + (SELECT generation FROM control_state WHERE singleton=TRUE), + (SELECT count(*) FROM tenants WHERE id=$1 AND status='active'), + (SELECT count(*) FROM projects WHERE tenant_id=$1 AND status='active'), + (SELECT count(*) FROM api_keys WHERE tenant_id=$1 AND status='active'), + 0, + (SELECT count(*) FROM models WHERE enabled=TRUE)`, tenantID, + ).Scan(&result.Generation, &result.Tenants, &result.Projects, &result.APIKeys, &result.Providers, &result.Models) + if err != nil { + return Overview{}, fmt.Errorf("query tenant overview: %w", err) + } + return result, nil +} + +func (s *Store) ResourceTenantID(ctx context.Context, resource, id string) (string, error) { + var query string + switch resource { + case "project": + query = `SELECT tenant_id::text FROM projects WHERE id=$1` + case "api_key": + query = `SELECT tenant_id::text FROM api_keys WHERE id=$1` + case "console_user": + query = `SELECT COALESCE(tenant_id::text,'') FROM console_users WHERE id=$1` + default: + return "", ErrNotFound + } + var tenantID string + if err := s.db.QueryRow(ctx, query, id).Scan(&tenantID); err != nil { + return "", ErrNotFound + } + return tenantID, nil +} + func (s *Store) ListProviders(ctx context.Context) ([]Provider, error) { rows, err := s.db.Query(ctx, ` SELECT p.id::text, p.name, p.protocol, p.base_url, p.enabled, count(r.id), p.created_at @@ -101,7 +172,11 @@ func (s *Store) ListProviders(ctx context.Context) ([]Provider, error) { } func (s *Store) ListModels(ctx context.Context) ([]Model, error) { - rows, err := s.db.Query(ctx, `SELECT id::text, public_id, owned_by, enabled, created_at FROM models ORDER BY public_id`) + rows, err := s.db.Query(ctx, ` + SELECT id::text, public_id, owned_by, input_price_micros_per_million, + output_price_micros_per_million, cache_read_price_micros_per_million, + cache_write_price_micros_per_million, enabled, created_at + FROM models ORDER BY public_id`) if err != nil { return nil, fmt.Errorf("query models: %w", err) } @@ -109,7 +184,9 @@ func (s *Store) ListModels(ctx context.Context) ([]Model, error) { positions := make(map[string]int) for rows.Next() { var item Model - if err := rows.Scan(&item.ID, &item.PublicID, &item.OwnedBy, &item.Enabled, &item.CreatedAt); err != nil { + if err := rows.Scan(&item.ID, &item.PublicID, &item.OwnedBy, &item.InputPriceMicrosPerMillion, + &item.OutputPriceMicrosPerMillion, &item.CacheReadPriceMicrosPerMillion, + &item.CacheWritePriceMicrosPerMillion, &item.Enabled, &item.CreatedAt); err != nil { rows.Close() return nil, fmt.Errorf("scan model: %w", err) } diff --git a/internal/controlplane/schema.sql b/internal/controlplane/schema.sql index 25af7c9..a518b1d 100644 --- a/internal/controlplane/schema.sql +++ b/internal/controlplane/schema.sql @@ -58,11 +58,20 @@ CREATE TABLE IF NOT EXISTS models ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), public_id TEXT NOT NULL UNIQUE, owned_by TEXT NOT NULL DEFAULT '', + input_price_micros_per_million BIGINT NOT NULL DEFAULT 0 CHECK (input_price_micros_per_million >= 0), + output_price_micros_per_million BIGINT NOT NULL DEFAULT 0 CHECK (output_price_micros_per_million >= 0), + cache_read_price_micros_per_million BIGINT NOT NULL DEFAULT 0 CHECK (cache_read_price_micros_per_million >= 0), + cache_write_price_micros_per_million BIGINT NOT NULL DEFAULT 0 CHECK (cache_write_price_micros_per_million >= 0), enabled BOOLEAN NOT NULL DEFAULT TRUE, created_at TIMESTAMPTZ NOT NULL DEFAULT now(), updated_at TIMESTAMPTZ NOT NULL DEFAULT now() ); +ALTER TABLE models ADD COLUMN IF NOT EXISTS input_price_micros_per_million BIGINT NOT NULL DEFAULT 0 CHECK (input_price_micros_per_million >= 0); +ALTER TABLE models ADD COLUMN IF NOT EXISTS output_price_micros_per_million BIGINT NOT NULL DEFAULT 0 CHECK (output_price_micros_per_million >= 0); +ALTER TABLE models ADD COLUMN IF NOT EXISTS cache_read_price_micros_per_million BIGINT NOT NULL DEFAULT 0 CHECK (cache_read_price_micros_per_million >= 0); +ALTER TABLE models ADD COLUMN IF NOT EXISTS cache_write_price_micros_per_million BIGINT NOT NULL DEFAULT 0 CHECK (cache_write_price_micros_per_million >= 0); + CREATE TABLE IF NOT EXISTS model_routes ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), model_id UUID NOT NULL REFERENCES models(id) ON DELETE CASCADE, @@ -80,3 +89,173 @@ CREATE INDEX IF NOT EXISTS api_keys_active_hash_idx ON api_keys (key_hash) WHERE CREATE INDEX IF NOT EXISTS projects_tenant_idx ON projects (tenant_id); CREATE INDEX IF NOT EXISTS model_routes_model_idx ON model_routes (model_id) WHERE enabled; CREATE INDEX IF NOT EXISTS model_routes_provider_idx ON model_routes (provider_id) WHERE enabled; + +CREATE TABLE IF NOT EXISTS tenant_wallets ( + tenant_id UUID PRIMARY KEY REFERENCES tenants(id) ON DELETE CASCADE, + currency TEXT NOT NULL CHECK (currency = lower(currency) AND length(currency) = 3), + balance_micros BIGINT NOT NULL DEFAULT 0 CHECK (balance_micros >= 0), + reserved_micros BIGINT NOT NULL DEFAULT 0 CHECK (reserved_micros >= 0 AND reserved_micros <= balance_micros), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +CREATE TABLE IF NOT EXISTS billing_reservations ( + request_id TEXT PRIMARY KEY, + tenant_id UUID NOT NULL REFERENCES tenants(id) ON DELETE CASCADE, + project_id UUID NOT NULL, + key_id UUID NOT NULL, + public_model TEXT NOT NULL, + currency TEXT NOT NULL, + reserved_micros BIGINT NOT NULL CHECK (reserved_micros >= 0), + status TEXT NOT NULL DEFAULT 'pending' CHECK (status IN ('pending', 'settled', 'released')), + input_price_micros_per_million BIGINT NOT NULL, + output_price_micros_per_million BIGINT NOT NULL, + cache_read_price_micros_per_million BIGINT NOT NULL, + cache_write_price_micros_per_million BIGINT NOT NULL, + actual_cost_micros BIGINT NOT NULL DEFAULT 0 CHECK (actual_cost_micros >= 0), + charged_micros BIGINT NOT NULL DEFAULT 0 CHECK (charged_micros >= 0), + uncollected_micros BIGINT NOT NULL DEFAULT 0 CHECK (uncollected_micros >= 0), + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + settled_at TIMESTAMPTZ, + FOREIGN KEY (project_id, tenant_id) REFERENCES projects(id, tenant_id) ON DELETE CASCADE +); + +CREATE TABLE IF NOT EXISTS usage_events ( + request_id TEXT PRIMARY KEY, + tenant_id UUID NOT NULL REFERENCES tenants(id) ON DELETE RESTRICT, + project_id UUID NOT NULL, + key_id UUID NOT NULL, + public_model TEXT NOT NULL, + provider_id TEXT, + upstream_model TEXT, + protocol TEXT NOT NULL, + stream BOOLEAN NOT NULL DEFAULT FALSE, + status_code INTEGER NOT NULL, + success BOOLEAN NOT NULL, + error_type TEXT NOT NULL DEFAULT '', + attempts INTEGER NOT NULL DEFAULT 0, + started_at TIMESTAMPTZ NOT NULL, + duration_ms BIGINT NOT NULL DEFAULT 0, + input_tokens BIGINT NOT NULL DEFAULT 0, + output_tokens BIGINT NOT NULL DEFAULT 0, + total_tokens BIGINT NOT NULL DEFAULT 0, + cache_creation_input_tokens BIGINT NOT NULL DEFAULT 0, + cache_read_input_tokens BIGINT NOT NULL DEFAULT 0, + cost_micros BIGINT NOT NULL DEFAULT 0, + charged_micros BIGINT NOT NULL DEFAULT 0, + uncollected_micros BIGINT NOT NULL DEFAULT 0, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + FOREIGN KEY (project_id, tenant_id) REFERENCES projects(id, tenant_id) ON DELETE RESTRICT +); + +-- Usage persistence is independent from billing. Older installations created this +-- foreign key, which prevented recording requests when prepaid billing was disabled. +ALTER TABLE usage_events DROP CONSTRAINT IF EXISTS usage_events_request_id_fkey; + +CREATE TABLE IF NOT EXISTS billing_ledger ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + tenant_id UUID NOT NULL REFERENCES tenants(id) ON DELETE RESTRICT, + project_id UUID, + currency TEXT NOT NULL, + amount_micros BIGINT NOT NULL, + balance_after_micros BIGINT NOT NULL CHECK (balance_after_micros >= 0), + kind TEXT NOT NULL CHECK (kind IN ('topup', 'usage', 'adjustment', 'refund', 'release')), + source_type TEXT NOT NULL, + source_id TEXT NOT NULL, + description TEXT NOT NULL DEFAULT '', + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + UNIQUE (source_type, source_id) +); + +CREATE TABLE IF NOT EXISTS topup_orders ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + tenant_id UUID NOT NULL REFERENCES tenants(id) ON DELETE RESTRICT, + amount_minor BIGINT NOT NULL CHECK (amount_minor > 0), + amount_micros BIGINT NOT NULL CHECK (amount_micros > 0), + currency TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'pending' CHECK (status IN ('pending', 'paid', 'failed', 'expired')), + stripe_session_id TEXT UNIQUE, + checkout_url TEXT, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + paid_at TIMESTAMPTZ +); + +CREATE TABLE IF NOT EXISTS stripe_webhook_events ( + event_id TEXT PRIMARY KEY, + event_type TEXT NOT NULL, + processed_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +CREATE TABLE IF NOT EXISTS console_users ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + tenant_id UUID REFERENCES tenants(id) ON DELETE CASCADE, + email TEXT NOT NULL, + display_name TEXT NOT NULL, + role TEXT NOT NULL CHECK (role IN ( + 'platform_admin', 'platform_viewer', 'tenant_admin', + 'tenant_billing', 'tenant_developer', 'tenant_viewer' + )), + token_prefix TEXT NOT NULL, + token_hash BYTEA NOT NULL UNIQUE CHECK (octet_length(token_hash) = 32), + status TEXT NOT NULL DEFAULT 'active' CHECK (status IN ('active', 'revoked')), + last_used_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + revoked_at TIMESTAMPTZ, + CHECK ((role LIKE 'platform_%' AND tenant_id IS NULL) OR (role LIKE 'tenant_%' AND tenant_id IS NOT NULL)) +); +CREATE UNIQUE INDEX IF NOT EXISTS console_users_email_tenant_idx + ON console_users (lower(email), COALESCE(tenant_id, '00000000-0000-0000-0000-000000000000'::uuid)); + +CREATE TABLE IF NOT EXISTS project_limits ( + project_id UUID PRIMARY KEY, + tenant_id UUID NOT NULL, + requests_per_minute BIGINT NOT NULL DEFAULT 0 CHECK (requests_per_minute >= 0), + tokens_per_minute BIGINT NOT NULL DEFAULT 0 CHECK (tokens_per_minute >= 0), + concurrent_requests BIGINT NOT NULL DEFAULT 0 CHECK (concurrent_requests >= 0), + monthly_spend_micros BIGINT NOT NULL DEFAULT 0 CHECK (monthly_spend_micros >= 0), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + FOREIGN KEY (project_id, tenant_id) REFERENCES projects(id, tenant_id) ON DELETE CASCADE +); + +CREATE TABLE IF NOT EXISTS usage_monthly_rollups ( + period_start DATE NOT NULL, + tenant_id UUID NOT NULL REFERENCES tenants(id) ON DELETE RESTRICT, + project_id UUID NOT NULL, + request_count BIGINT NOT NULL DEFAULT 0, + successful_requests BIGINT NOT NULL DEFAULT 0, + input_tokens BIGINT NOT NULL DEFAULT 0, + output_tokens BIGINT NOT NULL DEFAULT 0, + total_tokens BIGINT NOT NULL DEFAULT 0, + cost_micros BIGINT NOT NULL DEFAULT 0, + charged_micros BIGINT NOT NULL DEFAULT 0, + uncollected_micros BIGINT NOT NULL DEFAULT 0, + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + PRIMARY KEY (project_id, period_start), + FOREIGN KEY (project_id, tenant_id) REFERENCES projects(id, tenant_id) ON DELETE RESTRICT +); + +CREATE TABLE IF NOT EXISTS audit_logs ( + id BIGSERIAL PRIMARY KEY, + actor_id UUID REFERENCES console_users(id) ON DELETE SET NULL, + actor_type TEXT NOT NULL CHECK (actor_type IN ('bootstrap', 'console_user')), + actor_role TEXT NOT NULL, + tenant_id UUID REFERENCES tenants(id) ON DELETE SET NULL, + request_id TEXT NOT NULL, + method TEXT NOT NULL, + path TEXT NOT NULL, + action TEXT NOT NULL, + status_code INTEGER NOT NULL, + remote_ip INET, + user_agent TEXT NOT NULL DEFAULT '', + created_at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +CREATE INDEX IF NOT EXISTS billing_ledger_tenant_idx ON billing_ledger (tenant_id, created_at DESC); +CREATE INDEX IF NOT EXISTS usage_events_tenant_idx ON usage_events (tenant_id, created_at DESC); +CREATE INDEX IF NOT EXISTS usage_events_project_idx ON usage_events (project_id, created_at DESC); +CREATE INDEX IF NOT EXISTS usage_events_model_idx ON usage_events (public_model, created_at DESC); +CREATE INDEX IF NOT EXISTS billing_reservations_pending_idx ON billing_reservations (status, created_at) WHERE status = 'pending'; +CREATE INDEX IF NOT EXISTS billing_reservations_project_pending_idx ON billing_reservations (project_id, created_at) WHERE status = 'pending'; +CREATE INDEX IF NOT EXISTS console_users_tenant_idx ON console_users (tenant_id, created_at DESC); +CREATE INDEX IF NOT EXISTS audit_logs_created_idx ON audit_logs (created_at DESC); +CREATE INDEX IF NOT EXISTS audit_logs_tenant_idx ON audit_logs (tenant_id, created_at DESC); diff --git a/internal/controlplane/snapshot.go b/internal/controlplane/snapshot.go index c8931ab..bc04e7c 100644 --- a/internal/controlplane/snapshot.go +++ b/internal/controlplane/snapshot.go @@ -37,12 +37,40 @@ func (s *Store) LoadSnapshot(ctx context.Context) (Snapshot, error) { if err != nil { return Snapshot{}, err } + result.Limits, err = loadLimitPolicies(ctx, tx) + if err != nil { + return Snapshot{}, err + } if err := tx.Commit(ctx); err != nil { return Snapshot{}, fmt.Errorf("commit snapshot transaction: %w", err) } return result, nil } +func loadLimitPolicies(ctx context.Context, tx pgx.Tx) ([]domain.LimitPolicy, error) { + rows, err := tx.Query(ctx, ` + SELECT tenant_id::text, project_id::text, requests_per_minute, tokens_per_minute, + concurrent_requests, monthly_spend_micros + FROM project_limits`) + if err != nil { + return nil, fmt.Errorf("query project limits: %w", err) + } + defer rows.Close() + result := make([]domain.LimitPolicy, 0) + for rows.Next() { + var policy domain.LimitPolicy + if err := rows.Scan(&policy.TenantID, &policy.ProjectID, &policy.RequestsPerMinute, + &policy.TokensPerMinute, &policy.Concurrent, &policy.MonthlySpendMicros); err != nil { + return nil, fmt.Errorf("scan project limit: %w", err) + } + result = append(result, policy) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("read project limits: %w", err) + } + return result, nil +} + func (s *Store) loadProviders(ctx context.Context, tx pgx.Tx) (map[string]domain.Provider, error) { rows, err := tx.Query(ctx, ` SELECT id::text, name, protocol, base_url, api_key_ciphertext @@ -74,7 +102,9 @@ func (s *Store) loadProviders(ctx context.Context, tx pgx.Tx) (map[string]domain func loadModels(ctx context.Context, tx pgx.Tx, providers map[string]domain.Provider) ([]domain.Model, error) { rows, err := tx.Query(ctx, ` - SELECT m.public_id, m.owned_by, r.provider_id::text, r.upstream_model, r.priority, r.weight + SELECT m.public_id, m.owned_by, m.input_price_micros_per_million, m.output_price_micros_per_million, + m.cache_read_price_micros_per_million, m.cache_write_price_micros_per_million, + r.provider_id::text, r.upstream_model, r.priority, r.weight FROM models m JOIN model_routes r ON r.model_id = m.id AND r.enabled = TRUE JOIN providers p ON p.id = r.provider_id AND p.enabled = TRUE @@ -88,8 +118,9 @@ func loadModels(ctx context.Context, tx pgx.Tx, providers map[string]domain.Prov index := make(map[string]int) for rows.Next() { var publicID, ownedBy, providerID, upstreamModel string + var inputPrice, outputPrice, cacheReadPrice, cacheWritePrice int64 var priority, weight int - if err := rows.Scan(&publicID, &ownedBy, &providerID, &upstreamModel, &priority, &weight); err != nil { + if err := rows.Scan(&publicID, &ownedBy, &inputPrice, &outputPrice, &cacheReadPrice, &cacheWritePrice, &providerID, &upstreamModel, &priority, &weight); err != nil { return nil, fmt.Errorf("scan model route: %w", err) } provider, ok := providers[providerID] @@ -100,7 +131,11 @@ func loadModels(ctx context.Context, tx pgx.Tx, providers map[string]domain.Prov if !exists { position = len(models) index[publicID] = position - models = append(models, domain.Model{ID: publicID, OwnedBy: ownedBy}) + models = append(models, domain.Model{ + ID: publicID, OwnedBy: ownedBy, + InputPriceMicrosPerMillion: inputPrice, OutputPriceMicrosPerMillion: outputPrice, + CacheReadPriceMicrosPerMillion: cacheReadPrice, CacheWritePriceMicrosPerMillion: cacheWritePrice, + }) } models[position].Routes = append(models[position].Routes, domain.Route{ Provider: provider, UpstreamModel: upstreamModel, Priority: priority, Weight: weight, diff --git a/internal/controlplane/types.go b/internal/controlplane/types.go index 24f2843..7dfb734 100644 --- a/internal/controlplane/types.go +++ b/internal/controlplane/types.go @@ -62,30 +62,38 @@ type Route struct { } type Model struct { - ID string `json:"id"` - PublicID string `json:"public_id"` - OwnedBy string `json:"owned_by"` - Enabled bool `json:"enabled"` - Routes []Route `json:"routes"` - CreatedAt time.Time `json:"created_at"` + ID string `json:"id"` + PublicID string `json:"public_id"` + OwnedBy string `json:"owned_by"` + InputPriceMicrosPerMillion int64 `json:"input_price_micros_per_million"` + OutputPriceMicrosPerMillion int64 `json:"output_price_micros_per_million"` + CacheReadPriceMicrosPerMillion int64 `json:"cache_read_price_micros_per_million"` + CacheWritePriceMicrosPerMillion int64 `json:"cache_write_price_micros_per_million"` + Enabled bool `json:"enabled"` + Routes []Route `json:"routes"` + CreatedAt time.Time `json:"created_at"` } type Overview struct { - Generation int64 `json:"generation"` - RuntimeGeneration int64 `json:"runtime_generation"` - RedisConfigured bool `json:"redis_configured"` - RedisConnected bool `json:"redis_connected"` - Tenants int64 `json:"tenants"` - Projects int64 `json:"projects"` - APIKeys int64 `json:"api_keys"` - Providers int64 `json:"providers"` - Models int64 `json:"models"` + Generation int64 `json:"generation"` + RuntimeGeneration int64 `json:"runtime_generation"` + RedisConfigured bool `json:"redis_configured"` + RedisConnected bool `json:"redis_connected"` + Tenants int64 `json:"tenants"` + Projects int64 `json:"projects"` + APIKeys int64 `json:"api_keys"` + Providers int64 `json:"providers"` + Models int64 `json:"models"` + BillingEnabled bool `json:"billing_enabled"` + StripeEnabled bool `json:"stripe_enabled"` + BillingCurrency string `json:"billing_currency,omitempty"` } type Snapshot struct { Generation int64 Models []domain.Model APIKeys []auth.HashedKeyRecord + Limits []domain.LimitPolicy } type ChangeEvent struct { @@ -132,7 +140,126 @@ type RouteInput struct { } type CreateModelInput struct { - PublicID string `json:"public_id"` - OwnedBy string `json:"owned_by"` - Routes []RouteInput `json:"routes"` + PublicID string `json:"public_id"` + OwnedBy string `json:"owned_by"` + InputPriceMicrosPerMillion int64 `json:"input_price_micros_per_million"` + OutputPriceMicrosPerMillion int64 `json:"output_price_micros_per_million"` + CacheReadPriceMicrosPerMillion int64 `json:"cache_read_price_micros_per_million"` + CacheWritePriceMicrosPerMillion int64 `json:"cache_write_price_micros_per_million"` + Routes []RouteInput `json:"routes"` +} + +type ConsoleActor struct { + ID string `json:"id,omitempty"` + TenantID string `json:"tenant_id,omitempty"` + Email string `json:"email"` + DisplayName string `json:"display_name"` + Role string `json:"role"` + Bootstrap bool `json:"bootstrap"` +} + +type ConsoleUser struct { + ID string `json:"id"` + TenantID string `json:"tenant_id,omitempty"` + Email string `json:"email"` + DisplayName string `json:"display_name"` + Role string `json:"role"` + TokenPrefix string `json:"token_prefix"` + Status string `json:"status"` + LastUsedAt *time.Time `json:"last_used_at,omitempty"` + CreatedAt time.Time `json:"created_at"` +} + +type CreatedConsoleUser struct { + ConsoleUser + Token string `json:"token"` +} + +type CreateConsoleUserInput struct { + TenantID string `json:"tenant_id"` + Email string `json:"email"` + DisplayName string `json:"display_name"` + Role string `json:"role"` +} + +type ProjectLimit struct { + domain.LimitPolicy + ProjectName string `json:"project_name"` + TenantName string `json:"tenant_name"` + UpdatedAt time.Time `json:"updated_at"` +} + +type SetProjectLimitInput struct { + RequestsPerMinute int64 `json:"requests_per_minute"` + TokensPerMinute int64 `json:"tokens_per_minute"` + ConcurrentRequests int64 `json:"concurrent_requests"` + MonthlySpendMicros int64 `json:"monthly_spend_micros"` +} + +type UsageRecord struct { + RequestID string `json:"request_id"` + TenantID string `json:"tenant_id"` + ProjectID string `json:"project_id"` + KeyID string `json:"key_id"` + PublicModel string `json:"public_model"` + ProviderID string `json:"provider_id,omitempty"` + UpstreamModel string `json:"upstream_model,omitempty"` + Protocol string `json:"protocol"` + Stream bool `json:"stream"` + StatusCode int `json:"status_code"` + Success bool `json:"success"` + ErrorType string `json:"error_type,omitempty"` + Attempts int `json:"attempts"` + StartedAt time.Time `json:"started_at"` + DurationMS int64 `json:"duration_ms"` + InputTokens int64 `json:"input_tokens"` + OutputTokens int64 `json:"output_tokens"` + TotalTokens int64 `json:"total_tokens"` + CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"` + CacheReadInputTokens int64 `json:"cache_read_input_tokens"` + CostMicros int64 `json:"cost_micros"` + ChargedMicros int64 `json:"charged_micros"` + UncollectedMicros int64 `json:"uncollected_micros"` +} + +type UsageSummary struct { + PeriodStart time.Time `json:"period_start"` + TenantID string `json:"tenant_id"` + ProjectID string `json:"project_id"` + ProjectName string `json:"project_name"` + RequestCount int64 `json:"request_count"` + SuccessfulRequests int64 `json:"successful_requests"` + InputTokens int64 `json:"input_tokens"` + OutputTokens int64 `json:"output_tokens"` + TotalTokens int64 `json:"total_tokens"` + CostMicros int64 `json:"cost_micros"` + ChargedMicros int64 `json:"charged_micros"` + UncollectedMicros int64 `json:"uncollected_micros"` +} + +type AuditLog struct { + ID int64 `json:"id"` + ActorID string `json:"actor_id,omitempty"` + ActorType string `json:"actor_type"` + ActorRole string `json:"actor_role"` + TenantID string `json:"tenant_id,omitempty"` + RequestID string `json:"request_id"` + Method string `json:"method"` + Path string `json:"path"` + Action string `json:"action"` + StatusCode int `json:"status_code"` + RemoteIP string `json:"remote_ip,omitempty"` + UserAgent string `json:"user_agent,omitempty"` + CreatedAt time.Time `json:"created_at"` +} + +type AuditInput struct { + Actor ConsoleActor + RequestID string + Method string + Path string + Action string + StatusCode int + RemoteIP string + UserAgent string } diff --git a/internal/controlplane/usage.go b/internal/controlplane/usage.go new file mode 100644 index 0000000..6c69a9f --- /dev/null +++ b/internal/controlplane/usage.go @@ -0,0 +1,151 @@ +package controlplane + +import ( + "context" + "fmt" + "strings" + "time" + + "aigw/internal/domain" + + "github.com/jackc/pgx/v5" +) + +type UsageQuery struct { + TenantID string + ProjectID string + Model string + Limit int +} + +func (s *Store) RecordUsage(ctx context.Context, event domain.UsageEvent) error { + tx, err := s.db.Begin(ctx) + if err != nil { + return fmt.Errorf("begin usage record: %w", err) + } + defer tx.Rollback(ctx) + command, err := tx.Exec(ctx, ` + INSERT INTO usage_events ( + request_id, tenant_id, project_id, key_id, public_model, provider_id, upstream_model, + protocol, stream, status_code, success, error_type, attempts, started_at, duration_ms, + input_tokens, output_tokens, total_tokens, cache_creation_input_tokens, cache_read_input_tokens, + cost_micros, charged_micros, uncollected_micros) + VALUES ($1,$2,$3,$4,$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,0,0,0) + ON CONFLICT (request_id) DO NOTHING`, event.RequestID, event.TenantID, event.ProjectID, event.KeyID, + event.PublicModel, event.ProviderID, event.UpstreamModel, string(event.Protocol), event.Stream, + event.StatusCode, event.Success, event.ErrorType, event.Attempts, event.StartedAt, event.DurationMS, + event.Usage.InputTokens, event.Usage.OutputTokens, event.Usage.TotalTokens, + event.Usage.CacheCreationInputTokens, event.Usage.CacheReadInputTokens) + if err != nil { + return fmt.Errorf("persist usage event: %w", err) + } + if command.RowsAffected() == 0 { + return tx.Commit(ctx) + } + if err := upsertUsageRollup(ctx, tx, event, 0, 0, 0); err != nil { + return err + } + if err := tx.Commit(ctx); err != nil { + return fmt.Errorf("commit usage record: %w", err) + } + return nil +} + +func upsertUsageRollup(ctx context.Context, tx pgx.Tx, event domain.UsageEvent, cost, charged, uncollected int64) error { + period := time.Date(event.StartedAt.UTC().Year(), event.StartedAt.UTC().Month(), 1, 0, 0, 0, 0, time.UTC) + _, err := tx.Exec(ctx, ` + INSERT INTO usage_monthly_rollups + (period_start, tenant_id, project_id, request_count, successful_requests, input_tokens, output_tokens, total_tokens, cost_micros, charged_micros, uncollected_micros) + VALUES ($1,$2,$3,1,$4,$5,$6,$7,$8,$9,$10) + ON CONFLICT (project_id, period_start) DO UPDATE SET request_count=usage_monthly_rollups.request_count+1, + successful_requests=usage_monthly_rollups.successful_requests+EXCLUDED.successful_requests, + input_tokens=usage_monthly_rollups.input_tokens+EXCLUDED.input_tokens, + output_tokens=usage_monthly_rollups.output_tokens+EXCLUDED.output_tokens, + total_tokens=usage_monthly_rollups.total_tokens+EXCLUDED.total_tokens, + cost_micros=usage_monthly_rollups.cost_micros+EXCLUDED.cost_micros, + charged_micros=usage_monthly_rollups.charged_micros+EXCLUDED.charged_micros, + uncollected_micros=usage_monthly_rollups.uncollected_micros+EXCLUDED.uncollected_micros, + updated_at=now()`, period, event.TenantID, event.ProjectID, boolInt(event.Success), + event.Usage.InputTokens, event.Usage.OutputTokens, event.Usage.TotalTokens, cost, charged, uncollected) + if err != nil { + return fmt.Errorf("update usage monthly rollup: %w", err) + } + return nil +} + +func boolInt(value bool) int { + if value { + return 1 + } + return 0 +} + +func (s *Store) ListUsage(ctx context.Context, query UsageQuery) ([]UsageRecord, error) { + limit := query.Limit + if limit < 1 || limit > 1000 { + limit = 200 + } + where := []string{"1=1"} + args := make([]any, 0, 5) + index := 1 + for _, item := range []struct{ value, clause string }{{query.TenantID, "tenant_id=$"}, {query.ProjectID, "project_id=$"}, {query.Model, "public_model=$"}} { + if strings.TrimSpace(item.value) != "" { + where = append(where, item.clause+fmt.Sprint(index)) + args = append(args, item.value) + index++ + } + } + args = append(args, limit) + rows, err := s.db.Query(ctx, `SELECT request_id, tenant_id::text, project_id::text, key_id::text, public_model, + COALESCE(provider_id,''), COALESCE(upstream_model,''), protocol, stream, status_code, success, error_type, + attempts, started_at, duration_ms, input_tokens, output_tokens, total_tokens, cache_creation_input_tokens, + cache_read_input_tokens, cost_micros, charged_micros, uncollected_micros FROM usage_events WHERE `+ + strings.Join(where, " AND ")+` ORDER BY created_at DESC LIMIT $`+fmt.Sprint(index), args...) + if err != nil { + return nil, fmt.Errorf("query usage events: %w", err) + } + defer rows.Close() + result := make([]UsageRecord, 0) + for rows.Next() { + var item UsageRecord + if err := rows.Scan(&item.RequestID, &item.TenantID, &item.ProjectID, &item.KeyID, &item.PublicModel, &item.ProviderID, &item.UpstreamModel, + &item.Protocol, &item.Stream, &item.StatusCode, &item.Success, &item.ErrorType, &item.Attempts, &item.StartedAt, &item.DurationMS, + &item.InputTokens, &item.OutputTokens, &item.TotalTokens, &item.CacheCreationInputTokens, &item.CacheReadInputTokens, + &item.CostMicros, &item.ChargedMicros, &item.UncollectedMicros); err != nil { + return nil, fmt.Errorf("scan usage event: %w", err) + } + result = append(result, item) + } + return result, rows.Err() +} + +func (s *Store) UsageSummary(ctx context.Context, tenantID, projectID string) ([]UsageSummary, error) { + where := []string{"1=1"} + args := make([]any, 0, 2) + index := 1 + for _, item := range []struct{ value, clause string }{{tenantID, "r.tenant_id=$"}, {projectID, "r.project_id=$"}} { + if strings.TrimSpace(item.value) != "" { + where = append(where, item.clause+fmt.Sprint(index)) + args = append(args, item.value) + index++ + } + } + rows, err := s.db.Query(ctx, `SELECT r.period_start, r.tenant_id::text, r.project_id::text, p.name, + r.request_count, r.successful_requests, r.input_tokens, r.output_tokens, r.total_tokens, + r.cost_micros, r.charged_micros, r.uncollected_micros FROM usage_monthly_rollups r JOIN projects p ON p.id=r.project_id WHERE `+ + strings.Join(where, " AND ")+` ORDER BY r.period_start DESC, p.name`, args...) + if err != nil { + return nil, fmt.Errorf("query usage summary: %w", err) + } + defer rows.Close() + result := make([]UsageSummary, 0) + for rows.Next() { + var item UsageSummary + if err := rows.Scan(&item.PeriodStart, &item.TenantID, &item.ProjectID, &item.ProjectName, &item.RequestCount, &item.SuccessfulRequests, + &item.InputTokens, &item.OutputTokens, &item.TotalTokens, &item.CostMicros, &item.ChargedMicros, &item.UncollectedMicros); err != nil { + return nil, fmt.Errorf("scan usage summary: %w", err) + } + result = append(result, item) + } + return result, rows.Err() +} diff --git a/internal/domain/types.go b/internal/domain/types.go index e2586ea..b1decc9 100644 --- a/internal/domain/types.go +++ b/internal/domain/types.go @@ -31,9 +31,13 @@ type Route struct { } type Model struct { - ID string - OwnedBy string - Routes []Route + ID string + OwnedBy string + InputPriceMicrosPerMillion int64 + OutputPriceMicrosPerMillion int64 + CacheReadPriceMicrosPerMillion int64 + CacheWritePriceMicrosPerMillion int64 + Routes []Route } type Usage struct { @@ -62,3 +66,14 @@ type UsageEvent struct { DurationMS int64 `json:"duration_ms"` Usage Usage `json:"usage"` } + +// LimitPolicy is the immutable runtime view of a project's commercial limits. +// Zero values disable the corresponding limit. +type LimitPolicy struct { + TenantID string `json:"tenant_id"` + ProjectID string `json:"project_id"` + RequestsPerMinute int64 `json:"requests_per_minute"` + TokensPerMinute int64 `json:"tokens_per_minute"` + Concurrent int64 `json:"concurrent_requests"` + MonthlySpendMicros int64 `json:"monthly_spend_micros"` +} diff --git a/internal/httpapi/api.go b/internal/httpapi/api.go index ed7939f..becfbf3 100644 --- a/internal/httpapi/api.go +++ b/internal/httpapi/api.go @@ -10,13 +10,16 @@ import ( "io" "log/slog" "net/http" + "runtime/debug" "strings" "time" "aigw/internal/apierror" "aigw/internal/auth" + "aigw/internal/billing" "aigw/internal/catalog" "aigw/internal/domain" + "aigw/internal/limits" "aigw/internal/provider" "aigw/internal/routing" "aigw/internal/telemetry" @@ -31,6 +34,11 @@ type API struct { router *routing.Router forwarder *provider.Forwarder usageSink telemetry.UsageSink + billingMeter billing.Meter + limiter *limits.Limiter + usageRecorder interface { + RecordUsage(context.Context, domain.UsageEvent) error + } metrics *telemetry.Metrics logger *slog.Logger maxBodyBytes int64 @@ -43,6 +51,11 @@ type Options struct { Router *routing.Router Forwarder *provider.Forwarder UsageSink telemetry.UsageSink + BillingMeter billing.Meter + Limiter *limits.Limiter + UsageRecorder interface { + RecordUsage(context.Context, domain.UsageEvent) error + } Metrics *telemetry.Metrics Logger *slog.Logger MaxBodyBytes int64 @@ -56,6 +69,9 @@ func New(options Options) *API { router: options.Router, forwarder: options.Forwarder, usageSink: options.UsageSink, + billingMeter: options.BillingMeter, + limiter: options.Limiter, + usageRecorder: options.UsageRecorder, metrics: options.Metrics, logger: options.Logger, maxBodyBytes: options.MaxBodyBytes, @@ -134,6 +150,25 @@ func (a *API) serveInference(w http.ResponseWriter, r *http.Request, protocol do apierror.Write(w, apierror.Error{Status: http.StatusBadRequest, Type: "invalid_params", Message: "Parameter model is required and the body must be valid JSON"}, requestID) return } + if a.limiter != nil { + lease, limitErr := a.limiter.Acquire(r.Context(), principal, body) + if limitErr != nil { + w.Header().Set("Retry-After", "60") + typeName, message := "rate_limit_exceeded", "Project request rate limit exceeded" + switch { + case errors.Is(limitErr, limits.ErrTokensExceeded): + typeName, message = "token_rate_limit_exceeded", "Project token rate limit exceeded" + case errors.Is(limitErr, limits.ErrConcurrencyLimit): + typeName, message = "concurrency_limit_exceeded", "Project concurrent request limit exceeded" + } + apierror.Write(w, apierror.Error{Status: http.StatusTooManyRequests, Type: typeName, Message: message}, requestID) + a.recordUsageOnly(r, domain.UsageEvent{RequestID: requestID, KeyID: principal.KeyID, TenantID: principal.TenantID, ProjectID: principal.ProjectID, + PublicModel: envelope.Model, Protocol: protocol, Stream: envelope.Stream, StatusCode: http.StatusTooManyRequests, Success: false, + ErrorType: typeName, StartedAt: startedAt, DurationMS: time.Since(startedAt).Milliseconds()}) + return + } + defer lease.Release() + } routes, err := a.router.Plan(envelope.Model, protocol) if err != nil { @@ -144,6 +179,41 @@ func (a *API) serveInference(w http.ResponseWriter, r *http.Request, protocol do } return } + if a.billingMeter != nil { + model, modelErr := a.catalog.Model(envelope.Model) + if modelErr != nil { + apierror.Write(w, apierror.Error{Status: http.StatusNotFound, Type: "invalid_model", Message: "Model does not exist"}, requestID) + return + } + policy := domain.LimitPolicy{} + if a.limiter != nil { + policy, _ = a.limiter.Policy(principal.ProjectID) + } + if err := a.billingMeter.Authorize(r.Context(), billing.Authorization{ + RequestID: requestID, Principal: principal, Model: model, Body: body, Policy: policy, + }); err != nil { + if errors.Is(err, billing.ErrInsufficientBalance) { + apierror.Write(w, apierror.Error{Status: http.StatusPaymentRequired, Type: "insufficient_balance", Message: "Account balance is insufficient"}, requestID) + a.recordUsageOnly(r, domain.UsageEvent{RequestID: requestID, KeyID: principal.KeyID, TenantID: principal.TenantID, ProjectID: principal.ProjectID, + PublicModel: envelope.Model, Protocol: protocol, Stream: envelope.Stream, StatusCode: http.StatusPaymentRequired, Success: false, + ErrorType: "insufficient_balance", StartedAt: startedAt, DurationMS: time.Since(startedAt).Milliseconds()}) + return + } + if errors.Is(err, billing.ErrQuotaExceeded) { + apierror.Write(w, apierror.Error{Status: http.StatusTooManyRequests, Type: "monthly_quota_exceeded", Message: "Project monthly spend quota exceeded"}, requestID) + a.recordUsageOnly(r, domain.UsageEvent{RequestID: requestID, KeyID: principal.KeyID, TenantID: principal.TenantID, ProjectID: principal.ProjectID, + PublicModel: envelope.Model, Protocol: protocol, Stream: envelope.Stream, StatusCode: http.StatusTooManyRequests, Success: false, + ErrorType: "monthly_quota_exceeded", StartedAt: startedAt, DurationMS: time.Since(startedAt).Milliseconds()}) + return + } + a.logger.Error("billing_authorization_failed", "request_id", requestID, "tenant_id", principal.TenantID, "error", err) + apierror.Write(w, apierror.Error{Status: http.StatusServiceUnavailable, Type: "billing_unavailable", Message: "Billing service is temporarily unavailable"}, requestID) + a.recordUsageOnly(r, domain.UsageEvent{RequestID: requestID, KeyID: principal.KeyID, TenantID: principal.TenantID, ProjectID: principal.ProjectID, + PublicModel: envelope.Model, Protocol: protocol, Stream: envelope.Stream, StatusCode: http.StatusServiceUnavailable, Success: false, + ErrorType: "billing_unavailable", StartedAt: startedAt, DurationMS: time.Since(startedAt).Milliseconds()}) + return + } + } result, err := a.forwarder.Forward(r.Context(), protocol, requestID, body, r.Header, routes) if err != nil { @@ -155,7 +225,7 @@ func (a *API) serveInference(w http.ResponseWriter, r *http.Request, protocol do } else { apierror.Write(w, apierror.Error{Status: status, Type: errorType, Message: "No upstream provider is currently available"}, requestID) } - a.publishUsage(domain.UsageEvent{ + a.finishUsage(r, domain.UsageEvent{ RequestID: requestID, KeyID: principal.KeyID, TenantID: principal.TenantID, ProjectID: principal.ProjectID, PublicModel: envelope.Model, Protocol: protocol, Stream: envelope.Stream, StatusCode: status, Success: false, ErrorType: errorType, Attempts: result.Attempts, StartedAt: startedAt, DurationMS: time.Since(startedAt).Milliseconds(), @@ -167,7 +237,7 @@ func (a *API) serveInference(w http.ResponseWriter, r *http.Request, protocol do if result.Response.StatusCode >= 400 { gatewayError := normalizeProviderError(result.Response) apierror.Write(w, gatewayError, requestID) - a.publishUsage(domain.UsageEvent{ + a.finishUsage(r, domain.UsageEvent{ RequestID: requestID, KeyID: principal.KeyID, TenantID: principal.TenantID, ProjectID: principal.ProjectID, PublicModel: envelope.Model, ProviderID: result.Route.Provider.ID, UpstreamModel: result.Route.UpstreamModel, Protocol: protocol, Stream: envelope.Stream, StatusCode: gatewayError.Status, Success: false, @@ -192,7 +262,7 @@ func (a *API) serveInference(w http.ResponseWriter, r *http.Request, protocol do if copyErr != nil { errorType = "stream_interrupted" } - a.publishUsage(domain.UsageEvent{ + a.finishUsage(r, domain.UsageEvent{ RequestID: requestID, KeyID: principal.KeyID, TenantID: principal.TenantID, ProjectID: principal.ProjectID, PublicModel: envelope.Model, ProviderID: result.Route.Provider.ID, UpstreamModel: result.Route.UpstreamModel, Protocol: protocol, Stream: stream, StatusCode: result.Response.StatusCode, Success: success, @@ -254,6 +324,41 @@ func (a *API) publishUsage(event domain.UsageEvent) { } } +func (a *API) finishUsage(r *http.Request, event domain.UsageEvent) { + settled := false + if a.billingMeter != nil { + ctx, cancel := context.WithTimeout(context.WithoutCancel(r.Context()), 5*time.Second) + err := a.billingMeter.Settle(ctx, event) + cancel() + if err != nil { + a.logger.Error("billing_settlement_failed", "request_id", event.RequestID, "tenant_id", event.TenantID, "error", err) + } else { + settled = true + } + } + if a.usageRecorder != nil && !settled { + ctx, cancel := context.WithTimeout(context.WithoutCancel(r.Context()), 5*time.Second) + err := a.usageRecorder.RecordUsage(ctx, event) + cancel() + if err != nil { + a.logger.Error("usage_persistence_failed", "request_id", event.RequestID, "tenant_id", event.TenantID, "error", err) + } + } + a.publishUsage(event) +} + +func (a *API) recordUsageOnly(r *http.Request, event domain.UsageEvent) { + if a.usageRecorder != nil { + ctx, cancel := context.WithTimeout(context.WithoutCancel(r.Context()), 5*time.Second) + err := a.usageRecorder.RecordUsage(ctx, event) + cancel() + if err != nil { + a.logger.Error("usage_persistence_failed", "request_id", event.RequestID, "tenant_id", event.TenantID, "error", err) + } + } + a.publishUsage(event) +} + func (a *API) withRequestID(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { id := newRequestID() @@ -266,7 +371,7 @@ func (a *API) recoverPanics(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { defer func() { if recovered := recover(); recovered != nil { - a.logger.Error("request_panic", "request_id", requestIDFrom(r.Context()), "error", recovered) + a.logger.Error("request_panic", "request_id", requestIDFrom(r.Context()), "error", recovered, "stack", string(debug.Stack())) apierror.Write(w, apierror.Error{Status: http.StatusInternalServerError, Type: "internal_server_error", Message: "Internal server error"}, requestIDFrom(r.Context())) } }() diff --git a/internal/httpapi/api_test.go b/internal/httpapi/api_test.go index 66a344b..014b99b 100644 --- a/internal/httpapi/api_test.go +++ b/internal/httpapi/api_test.go @@ -3,6 +3,7 @@ package httpapi import ( "bufio" "bytes" + "context" "encoding/json" "io" "log/slog" @@ -14,6 +15,7 @@ import ( "time" "aigw/internal/auth" + "aigw/internal/billing" "aigw/internal/catalog" "aigw/internal/config" "aigw/internal/domain" @@ -26,6 +28,20 @@ type captureUsageSink struct { events chan domain.UsageEvent } +type fakeBillingMeter struct { + authorizeErr error + settled chan domain.UsageEvent +} + +func (m *fakeBillingMeter) Authorize(context.Context, billing.Authorization) error { + return m.authorizeErr +} + +func (m *fakeBillingMeter) Settle(_ context.Context, event domain.UsageEvent) error { + m.settled <- event + return nil +} + func (s *captureUsageSink) Publish(event domain.UsageEvent) { s.events <- event } @@ -178,7 +194,73 @@ func TestAnthropicHeadersAndPath(t *testing.T) { } } +func TestInsufficientBalanceRejectsBeforeCallingUpstream(t *testing.T) { + var calls atomic.Int64 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls.Add(1) + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + meter := &fakeBillingMeter{authorizeErr: billing.ErrInsufficientBalance, settled: make(chan domain.UsageEvent, 1)} + gateway, sink := newTestGatewayWithBilling(t, []config.ProviderConfig{{ + ID: "primary", Protocol: domain.ProtocolOpenAI, BaseURL: upstream.URL + "/v1", APIKey: "secret", + }}, []config.RouteConfig{{Provider: "primary", UpstreamModel: "model", Weight: 1}}, meter) + defer gateway.Close() + + response := postOpenAI(t, gateway.URL, false) + defer response.Body.Close() + if response.StatusCode != http.StatusPaymentRequired { + t.Fatalf("status = %d, want 402", response.StatusCode) + } + if calls.Load() != 0 { + t.Fatalf("upstream calls = %d, want 0", calls.Load()) + } + select { + case <-meter.settled: + t.Fatal("rejected request was settled") + default: + } + select { + case event := <-sink.events: + if event.StatusCode != http.StatusPaymentRequired || event.ErrorType != "insufficient_balance" { + t.Fatalf("unexpected rejection usage: %+v", event) + } + case <-time.After(time.Second): + t.Fatal("rejected request did not emit usage") + } +} + +func TestSuccessfulRequestIsSettled(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, `{"choices":[],"usage":{"prompt_tokens":2,"completion_tokens":3,"total_tokens":5}}`) + })) + defer upstream.Close() + meter := &fakeBillingMeter{settled: make(chan domain.UsageEvent, 1)} + gateway, _ := newTestGatewayWithBilling(t, []config.ProviderConfig{{ + ID: "primary", Protocol: domain.ProtocolOpenAI, BaseURL: upstream.URL + "/v1", APIKey: "secret", + }}, []config.RouteConfig{{Provider: "primary", UpstreamModel: "model", Weight: 1}}, meter) + defer gateway.Close() + + response := postOpenAI(t, gateway.URL, false) + defer response.Body.Close() + if response.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want 200", response.StatusCode) + } + select { + case event := <-meter.settled: + if event.Usage.TotalTokens != 5 { + t.Fatalf("settled usage = %+v", event.Usage) + } + case <-time.After(time.Second): + t.Fatal("request was not settled") + } +} + func newTestGateway(t *testing.T, providers []config.ProviderConfig, routes []config.RouteConfig) (*httptest.Server, *captureUsageSink) { + return newTestGatewayWithBilling(t, providers, routes, nil) +} + +func newTestGatewayWithBilling(t *testing.T, providers []config.ProviderConfig, routes []config.RouteConfig, meter billing.Meter) (*httptest.Server, *captureUsageSink) { t.Helper() authenticator, err := auth.NewStatic(`[{"key":"client-secret","key_id":"key-1","tenant_id":"tenant-1","project_id":"project-1","scopes":["inference"]}]`, false) if err != nil { @@ -201,6 +283,7 @@ func newTestGateway(t *testing.T, providers []config.ProviderConfig, routes []co Router: routing.New(modelCatalog), Forwarder: provider.New(cfg.UpstreamHTTP, metrics), UsageSink: sink, + BillingMeter: meter, Metrics: metrics, Logger: slog.New(slog.NewTextHandler(io.Discard, nil)), MaxBodyBytes: 1 << 20, diff --git a/internal/limits/limits.go b/internal/limits/limits.go new file mode 100644 index 0000000..6e0bf74 --- /dev/null +++ b/internal/limits/limits.go @@ -0,0 +1,316 @@ +package limits + +import ( + "context" + "encoding/json" + "errors" + "log/slog" + "math" + "strconv" + "strings" + "sync" + "sync/atomic" + "time" + + "aigw/internal/domain" + + "github.com/redis/go-redis/v9" +) + +var ( + ErrRequestsExceeded = errors.New("requests per minute limit exceeded") + ErrTokensExceeded = errors.New("tokens per minute limit exceeded") + ErrConcurrencyLimit = errors.New("concurrent request limit exceeded") +) + +type Lease interface{ Release() } + +type Limiter struct { + redis *redis.Client + prefix string + defaultMaxOutput int64 + logger *slog.Logger + policies atomic.Pointer[policySnapshot] + localMu sync.Mutex + local map[string]*localWindow + redisUp atomic.Bool + redisFailureSeen atomic.Bool + redisRetryAt atomic.Int64 +} + +type policySnapshot struct{ policies map[string]domain.LimitPolicy } + +type localWindow struct { + minute int64 + requests int64 + tokens int64 + concurrent int64 +} + +type localLease struct { + limiter *Limiter + project string + released atomic.Bool +} +type redisLease struct { + limiter *Limiter + key string + released atomic.Bool +} + +const ( + redisCommandTimeout = 250 * time.Millisecond + redisRetryCooldown = 5 * time.Second +) + +const acquireScript = ` +local req = tonumber(ARGV[1]) +local tok = tonumber(ARGV[2]) +local conc = tonumber(ARGV[3]) +local estimate = tonumber(ARGV[4]) +local ttl = tonumber(ARGV[5]) +local r = 0 +local t = 0 +local c = 0 +if req > 0 then r = redis.call('INCR', KEYS[1]); if r == 1 then redis.call('PEXPIRE', KEYS[1], ttl) end end +if tok > 0 then t = redis.call('INCRBY', KEYS[2], estimate); if t == estimate then redis.call('PEXPIRE', KEYS[2], ttl) end end +if conc > 0 then c = redis.call('INCR', KEYS[3]); redis.call('PEXPIRE', KEYS[3], 3600000) end +if (req > 0 and r > req) then if req > 0 then redis.call('DECR', KEYS[1]) end; if tok > 0 then redis.call('DECRBY', KEYS[2], estimate) end; if conc > 0 then redis.call('DECR', KEYS[3]) end; return {0,1} end +if (tok > 0 and t > tok) then if req > 0 then redis.call('DECR', KEYS[1]) end; if tok > 0 then redis.call('DECRBY', KEYS[2], estimate) end; if conc > 0 then redis.call('DECR', KEYS[3]) end; return {0,2} end +if (conc > 0 and c > conc) then if req > 0 then redis.call('DECR', KEYS[1]) end; if tok > 0 then redis.call('DECRBY', KEYS[2], estimate) end; if conc > 0 then redis.call('DECR', KEYS[3]) end; return {0,3} end +return {1,0} +` + +const releaseScript = ` +local value = tonumber(redis.call('GET', KEYS[1]) or '0') +if value > 0 then redis.call('DECR', KEYS[1]) end +return value +` + +func New(redisURL, prefix string, defaultMaxOutput int64, logger *slog.Logger) *Limiter { + if logger == nil { + logger = slog.Default() + } + if prefix == "" { + prefix = "aigw:limits" + } + l := &Limiter{prefix: strings.TrimRight(prefix, ":"), defaultMaxOutput: defaultMaxOutput, logger: logger, local: make(map[string]*localWindow)} + if strings.TrimSpace(redisURL) != "" { + if options, err := redis.ParseURL(redisURL); err == nil { + options.MaxRetries = -1 + options.DialerRetries = 1 + options.DialTimeout = redisCommandTimeout + options.ReadTimeout = redisCommandTimeout + options.WriteTimeout = redisCommandTimeout + options.PoolTimeout = redisCommandTimeout + l.redis = redis.NewClient(options) + } else { + logger.Warn("limits_redis_config_invalid", "error", err, "fallback", "local") + } + } + l.ReplacePolicies(nil) + return l +} + +func (l *Limiter) Close() error { + if l.redis == nil { + return nil + } + return l.redis.Close() +} + +func (l *Limiter) ReplacePolicies(policies []domain.LimitPolicy) { + copyMap := make(map[string]domain.LimitPolicy, len(policies)) + for _, policy := range policies { + if policy.ProjectID != "" { + copyMap[policy.ProjectID] = policy + } + } + l.policies.Store(&policySnapshot{policies: copyMap}) +} + +func (l *Limiter) Policy(projectID string) (domain.LimitPolicy, bool) { + snapshot := l.policies.Load() + if snapshot == nil { + return domain.LimitPolicy{}, false + } + policy, ok := snapshot.policies[projectID] + return policy, ok +} + +func (l *Limiter) Acquire(ctx context.Context, principal domain.Principal, body []byte) (Lease, error) { + policy, ok := l.Policy(principal.ProjectID) + if !ok || (policy.RequestsPerMinute == 0 && policy.TokensPerMinute == 0 && policy.Concurrent == 0) { + return noopLease{}, nil + } + estimate := EstimateTokens(body, l.defaultMaxOutput) + minute := time.Now().Unix() / 60 + if l.redis != nil && time.Now().UnixNano() >= l.redisRetryAt.Load() { + redisContext, cancel := context.WithTimeout(ctx, redisCommandTimeout) + lease, err := l.acquireRedis(redisContext, principal.ProjectID, minute, policy, estimate) + cancel() + if err == nil { + return lease, nil + } + if errors.Is(err, ErrRequestsExceeded) || errors.Is(err, ErrTokensExceeded) || errors.Is(err, ErrConcurrencyLimit) { + l.markRedis(true, nil) + return nil, err + } + l.markRedis(false, err) + } + return l.acquireLocal(principal.ProjectID, minute, policy, estimate) +} + +func EstimateTokens(body []byte, defaultMaxOutput int64) int64 { + input := int64((len(body) + 3) / 4) + if input < 1 { + input = 1 + } + maxOutput := defaultMaxOutput + var limits struct { + MaxTokens int64 `json:"max_tokens"` + MaxCompletionTokens int64 `json:"max_completion_tokens"` + MaxOutputTokens int64 `json:"max_output_tokens"` + } + if json.Unmarshal(body, &limits) == nil { + explicit := int64(0) + for _, value := range []int64{limits.MaxTokens, limits.MaxCompletionTokens, limits.MaxOutputTokens} { + if value > explicit { + explicit = value + } + } + if explicit > 0 { + maxOutput = explicit + } + } + if maxOutput < 0 { + maxOutput = 0 + } + if input > math.MaxInt64-maxOutput { + return math.MaxInt64 + } + return input + maxOutput +} + +func (l *Limiter) acquireRedis(ctx context.Context, project string, minute int64, policy domain.LimitPolicy, estimate int64) (Lease, error) { + base := l.prefix + ":" + project + ":" + strconv.FormatInt(minute, 10) + keys := []string{base + ":requests", base + ":tokens", l.prefix + ":" + project + ":concurrent"} + values, err := l.redis.Eval(ctx, acquireScript, keys, policy.RequestsPerMinute, policy.TokensPerMinute, policy.Concurrent, estimate, 125000).Result() + if err != nil { + return nil, err + } + items, ok := values.([]any) + if !ok || len(items) < 2 { + return nil, errors.New("invalid limits Redis response") + } + allowed, _ := toInt64(items[0]) + reason, _ := toInt64(items[1]) + if allowed == 0 { + switch reason { + case 1: + return nil, ErrRequestsExceeded + case 2: + return nil, ErrTokensExceeded + default: + return nil, ErrConcurrencyLimit + } + } + l.markRedis(true, nil) + return &redisLease{limiter: l, key: keys[2]}, nil +} + +func (l *Limiter) acquireLocal(project string, minute int64, policy domain.LimitPolicy, estimate int64) (Lease, error) { + l.localMu.Lock() + defer l.localMu.Unlock() + for key, value := range l.local { + if value.minute < minute-2 && value.concurrent == 0 { + delete(l.local, key) + } + } + window := l.local[project] + if window == nil { + window = &localWindow{minute: minute} + l.local[project] = window + } else if window.minute != minute { + window.minute = minute + window.requests = 0 + window.tokens = 0 + } + if policy.RequestsPerMinute > 0 && window.requests >= policy.RequestsPerMinute { + return nil, ErrRequestsExceeded + } + if policy.TokensPerMinute > 0 && (estimate > policy.TokensPerMinute || window.tokens > policy.TokensPerMinute-estimate) { + return nil, ErrTokensExceeded + } + if policy.Concurrent > 0 && window.concurrent >= policy.Concurrent { + return nil, ErrConcurrencyLimit + } + window.requests++ + window.tokens += estimate + window.concurrent++ + return &localLease{limiter: l, project: project}, nil +} + +func (l *Limiter) releaseLocal(project string) { + l.localMu.Lock() + defer l.localMu.Unlock() + if value := l.local[project]; value != nil && value.concurrent > 0 { + value.concurrent-- + } +} + +func (l *Limiter) releaseRedis(ctx context.Context, key string) { + if l.redis == nil { + return + } + if _, err := l.redis.Eval(ctx, releaseScript, []string{key}).Result(); err != nil { + l.markRedis(false, err) + } +} + +func (x *localLease) Release() { + if x.released.CompareAndSwap(false, true) { + x.limiter.releaseLocal(x.project) + } +} +func (x *redisLease) Release() { + if x.released.CompareAndSwap(false, true) { + ctx, cancel := context.WithTimeout(context.Background(), redisCommandTimeout) + defer cancel() + x.limiter.releaseRedis(ctx, x.key) + } +} + +type noopLease struct{} + +func (noopLease) Release() {} + +func (l *Limiter) markRedis(healthy bool, err error) { + l.redisUp.Store(healthy) + if healthy { + l.redisRetryAt.Store(0) + if l.redisFailureSeen.Swap(false) { + l.logger.Info("limits_redis_recovered") + } + return + } + l.redisRetryAt.Store(time.Now().Add(redisRetryCooldown).UnixNano()) + if l.redisFailureSeen.CompareAndSwap(false, true) { + l.logger.Warn("limits_redis_unavailable", "error", err, "fallback", "local", "retry_after", redisRetryCooldown) + } +} + +func toInt64(value any) (int64, bool) { + switch v := value.(type) { + case int64: + return v, true + case string: + n, e := strconv.ParseInt(v, 10, 64) + return n, e == nil + case []byte: + n, e := strconv.ParseInt(string(v), 10, 64) + return n, e == nil + default: + return 0, false + } +} diff --git a/internal/limits/limits_test.go b/internal/limits/limits_test.go new file mode 100644 index 0000000..78b346e --- /dev/null +++ b/internal/limits/limits_test.go @@ -0,0 +1,103 @@ +package limits + +import ( + "context" + "errors" + "testing" + "time" + + "aigw/internal/domain" +) + +func TestLocalRequestLimit(t *testing.T) { + limiter := New("", "test", 0, nil) + limiter.ReplacePolicies([]domain.LimitPolicy{{ProjectID: "project-1", RequestsPerMinute: 2}}) + principal := domain.Principal{ProjectID: "project-1"} + for i := 0; i < 2; i++ { + lease, err := limiter.Acquire(context.Background(), principal, []byte(`{}`)) + if err != nil { + t.Fatal(err) + } + lease.Release() + } + if _, err := limiter.Acquire(context.Background(), principal, []byte(`{}`)); !errors.Is(err, ErrRequestsExceeded) { + t.Fatalf("error = %v, want request limit", err) + } +} + +func TestLocalTokenAndConcurrencyLimits(t *testing.T) { + limiter := New("", "test", 0, nil) + body := []byte(`{"max_tokens":4}`) + estimate := EstimateTokens(body, 0) + limiter.ReplacePolicies([]domain.LimitPolicy{ + {ProjectID: "tokens", TokensPerMinute: estimate}, + {ProjectID: "concurrency", Concurrent: 1}, + }) + lease, err := limiter.Acquire(context.Background(), domain.Principal{ProjectID: "tokens"}, body) + if err != nil { + t.Fatal(err) + } + lease.Release() + if _, err := limiter.Acquire(context.Background(), domain.Principal{ProjectID: "tokens"}, body); !errors.Is(err, ErrTokensExceeded) { + t.Fatalf("error = %v, want token limit", err) + } + held, err := limiter.Acquire(context.Background(), domain.Principal{ProjectID: "concurrency"}, body) + if err != nil { + t.Fatal(err) + } + if _, err := limiter.Acquire(context.Background(), domain.Principal{ProjectID: "concurrency"}, body); !errors.Is(err, ErrConcurrencyLimit) { + t.Fatalf("error = %v, want concurrency limit", err) + } + held.Release() + retry, err := limiter.Acquire(context.Background(), domain.Principal{ProjectID: "concurrency"}, body) + if err != nil { + t.Fatalf("acquire after release: %v", err) + } + retry.Release() +} + +func TestEstimateTokensUsesLargestExplicitOutputLimit(t *testing.T) { + body := []byte(`{"max_tokens":10,"max_completion_tokens":25}`) + want := int64((len(body)+3)/4 + 25) + if got := EstimateTokens(body, 4); got != want { + t.Fatalf("EstimateTokens = %d, want %d", got, want) + } +} + +func TestEstimateTokensHonorsExplicitLimitBelowDefault(t *testing.T) { + body := []byte(`{"max_tokens":10}`) + want := int64((len(body)+3)/4 + 10) + if got := EstimateTokens(body, 4096); got != want { + t.Fatalf("EstimateTokens = %d, want %d", got, want) + } +} + +func TestRedisFailureFallsBackQuicklyAndOpensCircuit(t *testing.T) { + limiter := New("redis://127.0.0.1:1/0", "test", 0, nil) + defer limiter.Close() + limiter.ReplacePolicies([]domain.LimitPolicy{{ProjectID: "project-1", RequestsPerMinute: 2}}) + principal := domain.Principal{ProjectID: "project-1"} + + started := time.Now() + lease, err := limiter.Acquire(context.Background(), principal, []byte(`{}`)) + if err != nil { + t.Fatal(err) + } + lease.Release() + if elapsed := time.Since(started); elapsed > time.Second { + t.Fatalf("Redis fallback took %v, want under 1s", elapsed) + } + if limiter.redisRetryAt.Load() <= time.Now().UnixNano() { + t.Fatal("Redis failure did not open the retry circuit") + } + + started = time.Now() + lease, err = limiter.Acquire(context.Background(), principal, []byte(`{}`)) + if err != nil { + t.Fatal(err) + } + lease.Release() + if elapsed := time.Since(started); elapsed > 100*time.Millisecond { + t.Fatalf("open-circuit local fallback took %v, want under 100ms", elapsed) + } +} diff --git a/scripts/start-debug.sh b/scripts/start-debug.sh new file mode 100755 index 0000000..698ae68 --- /dev/null +++ b/scripts/start-debug.sh @@ -0,0 +1,72 @@ +#!/usr/bin/env bash + +set -euo pipefail + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +repo_dir="$(cd -- "$script_dir/.." && pwd)" +env_file="${AIGW_DEBUG_ENV_FILE:-$repo_dir/.env.debug}" + +log() { + printf '[aigw-debug] %s\n' "$*" +} + +fail() { + printf '[aigw-debug] error: %s\n' "$*" >&2 + exit 1 +} + +command -v docker >/dev/null 2>&1 || fail "docker is not installed" +docker info >/dev/null 2>&1 || fail "cannot access Docker; check that the daemon is running and your user belongs to the docker group" + +if [[ ! -f "$env_file" ]]; then + command -v openssl >/dev/null 2>&1 || fail "openssl is required to generate local credentials" + credential_key="$(openssl rand -base64 32 | tr -d '\n')" + admin_token="aigw-admin-$(openssl rand -hex 24)" + umask 077 + { + printf 'AIGW_DATABASE_URL=postgres://aigw:aigw@127.0.0.1:5432/aigw?sslmode=disable\n' + printf 'AIGW_REDIS_URL=redis://127.0.0.1:6379/0\n' + printf 'AIGW_CREDENTIAL_KEY=%s\n' "$credential_key" + printf 'AIGW_ADMIN_TOKEN=%s\n' "$admin_token" + printf 'AIGW_STRIPE_API_KEY=rk_test_replace_me\n' + printf 'AIGW_STRIPE_WEBHOOK_SECRET=whsec_replace_me\n' + } >"$env_file" + chmod 600 "$env_file" + log "created $env_file" +fi + +admin_token="$(sed -n 's/^AIGW_ADMIN_TOKEN=//p' "$env_file" | head -n 1)" +[[ -n "$admin_token" ]] || fail "AIGW_ADMIN_TOKEN is missing from $env_file" + +cd "$repo_dir" +if [[ "${AIGW_DEBUG_SKIP_BUILD:-0}" == "1" ]]; then + log "skipping image build (AIGW_DEBUG_SKIP_BUILD=1)" +else + build_network="${AIGW_DOCKER_BUILD_NETWORK:-host}" + log "building gateway image (network: $build_network)" + docker build --network="$build_network" -t aigw-debug:local . +fi + +log "starting PostgreSQL, Redis, and gateway" +if ! docker compose --env-file "$env_file" up -d --no-build --wait --wait-timeout 120; then + docker compose --env-file "$env_file" ps >&2 || true + gateway_logs="$(docker compose --env-file "$env_file" logs --no-color --tail=120 aigw 2>&1 || true)" + printf '%s\n' "$gateway_logs" >&2 + if [[ "$gateway_logs" == *"decrypt credential"* ]]; then + printf '[aigw-debug] the AIGW_CREDENTIAL_KEY in %s does not match credentials already stored in the PostgreSQL volume\n' "$env_file" >&2 + printf '[aigw-debug] restore the previous key, or explicitly remove the debug volumes if the stored control-plane data is disposable\n' >&2 + fi + fail "services did not become healthy" +fi + +log "services are ready" +printf '\nAdmin UI: http://127.0.0.1:8080/admin/\n' +printf 'Health: http://127.0.0.1:8080/readyz\n' +printf 'Admin token: %s\n' "$admin_token" +printf 'Environment: %s\n' "$env_file" +printf '\nLogs: docker compose --env-file %q logs -f aigw\n' "$env_file" +printf 'Stop: ./scripts/stop-debug.sh\n' + +if grep -q '^AIGW_STRIPE_API_KEY=rk_test_replace_me$' "$env_file" 2>/dev/null; then + printf '\nStripe uses placeholders. Replace the two Stripe values in %s before testing Checkout.\n' "$env_file" +fi diff --git a/scripts/stop-debug.sh b/scripts/stop-debug.sh new file mode 100755 index 0000000..8976d2c --- /dev/null +++ b/scripts/stop-debug.sh @@ -0,0 +1,27 @@ +#!/usr/bin/env bash + +set -euo pipefail + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +repo_dir="$(cd -- "$script_dir/.." && pwd)" +env_file="${AIGW_DEBUG_ENV_FILE:-$repo_dir/.env.debug}" + +command -v docker >/dev/null 2>&1 || { + printf '[aigw-debug] error: docker is not installed\n' >&2 + exit 1 +} + +if [[ -f "$env_file" ]]; then + compose=(docker compose --env-file "$env_file") +else + export AIGW_CREDENTIAL_KEY=debug-stop-placeholder + export AIGW_ADMIN_TOKEN=debug-stop-placeholder + export AIGW_STRIPE_API_KEY=debug-stop-placeholder + export AIGW_STRIPE_WEBHOOK_SECRET=debug-stop-placeholder + compose=(docker compose) +fi + +cd "$repo_dir" +printf '[aigw-debug] stopping services\n' +"${compose[@]}" down --remove-orphans +printf '[aigw-debug] stopped; PostgreSQL and Redis volumes were preserved\n' |