summaryrefslogtreecommitdiff
path: root/internal/controlplane
diff options
context:
space:
mode:
Diffstat (limited to 'internal/controlplane')
-rw-r--r--internal/controlplane/access.go156
-rw-r--r--internal/controlplane/access_test.go39
-rw-r--r--internal/controlplane/audit.go52
-rw-r--r--internal/controlplane/limits.go81
-rw-r--r--internal/controlplane/limits_integration_test.go62
-rw-r--r--internal/controlplane/manager.go20
-rw-r--r--internal/controlplane/manager_test.go20
-rw-r--r--internal/controlplane/mutations.go16
-rw-r--r--internal/controlplane/queries.go89
-rw-r--r--internal/controlplane/schema.sql179
-rw-r--r--internal/controlplane/snapshot.go41
-rw-r--r--internal/controlplane/types.go163
-rw-r--r--internal/controlplane/usage.go151
13 files changed, 1037 insertions, 32 deletions
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()
+}
generated by cgit v1.2.3 (git 2.46.0) at 2026-08-04 20:00:43 +0000