AI Gateway Go 0.10.0 源码快照 + 旗舰版需求规划报告
M0-M7 已完成:核心网关(身份/RBAC/TOTP/OIDC/SAML/Provider/配额/路由/内容策略/审计/定价)+ 资源市场(MCP/Skills/数字员工)。 含 22 个 PostgreSQL 迁移、管理端/门户端前端源码、OpenAPI 契约、部署 compose。 Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,144 @@
|
||||
package audit
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/identity"
|
||||
"aigateway.local/core/internal/platform/apiresponse"
|
||||
)
|
||||
|
||||
type AdminHTTPHandler struct {
|
||||
query *QueryService
|
||||
identity *identity.Service
|
||||
mux *http.ServeMux
|
||||
}
|
||||
|
||||
func NewAdminHTTPHandler(query *QueryService, identityService *identity.Service) *AdminHTTPHandler {
|
||||
h := &AdminHTTPHandler{query: query, identity: identityService, mux: http.NewServeMux()}
|
||||
h.mux.HandleFunc("GET /api/v1/admin/audit-events", h.listEvents)
|
||||
h.mux.HandleFunc("GET /api/v1/admin/usage/daily", h.listDailyUsage)
|
||||
return h
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
|
||||
h.mux.ServeHTTP(writer, request)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) listEvents(writer http.ResponseWriter, request *http.Request) {
|
||||
if !h.requirePermission(writer, request, identity.PermissionAuditRead) {
|
||||
return
|
||||
}
|
||||
now := time.Now().UTC()
|
||||
from, err := queryTime(request, "from", now.Add(-24*time.Hour))
|
||||
if err != nil {
|
||||
apiresponse.Error(writer, http.StatusBadRequest, "from 时间无效")
|
||||
return
|
||||
}
|
||||
to, err := queryTime(request, "to", now.Add(time.Second))
|
||||
if err != nil || !to.After(from) || to.Sub(from) > 366*24*time.Hour {
|
||||
apiresponse.Error(writer, http.StatusBadRequest, "审计查询时间范围无效或超过 366 天")
|
||||
return
|
||||
}
|
||||
limit := 50
|
||||
if value := request.URL.Query().Get("limit"); value != "" {
|
||||
limit, err = strconv.Atoi(value)
|
||||
if err != nil || limit < 1 || limit > 200 {
|
||||
apiresponse.Error(writer, http.StatusBadRequest, "limit 必须在 1 到 200 之间")
|
||||
return
|
||||
}
|
||||
}
|
||||
var before *time.Time
|
||||
if value := request.URL.Query().Get("before"); value != "" {
|
||||
parsed, parseErr := time.Parse(time.RFC3339Nano, value)
|
||||
if parseErr != nil {
|
||||
apiresponse.Error(writer, http.StatusBadRequest, "before 游标无效")
|
||||
return
|
||||
}
|
||||
before = &parsed
|
||||
}
|
||||
var status *int
|
||||
if value := request.URL.Query().Get("status"); value != "" {
|
||||
parsed, parseErr := strconv.Atoi(value)
|
||||
if parseErr != nil || parsed < 100 || parsed > 599 {
|
||||
apiresponse.Error(writer, http.StatusBadRequest, "status 无效")
|
||||
return
|
||||
}
|
||||
status = &parsed
|
||||
}
|
||||
items, err := h.query.ListEvents(request.Context(), EventFilter{
|
||||
From: from, To: to, Before: before, APIKeyID: strings.TrimSpace(request.URL.Query().Get("api_key_id")),
|
||||
Provider: strings.TrimSpace(request.URL.Query().Get("provider")), Model: strings.TrimSpace(request.URL.Query().Get("model")), StatusCode: status, Limit: limit,
|
||||
})
|
||||
if err != nil {
|
||||
apiresponse.Error(writer, http.StatusServiceUnavailable, "审计查询服务暂不可用")
|
||||
return
|
||||
}
|
||||
next := ""
|
||||
if len(items) == limit {
|
||||
next = items[len(items)-1].RecordedAt.Format(time.RFC3339Nano)
|
||||
}
|
||||
apiresponse.OK(writer, map[string]any{"items": items, "next_before": next})
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) listDailyUsage(writer http.ResponseWriter, request *http.Request) {
|
||||
if !h.requirePermission(writer, request, identity.PermissionUsageRead) {
|
||||
return
|
||||
}
|
||||
now := time.Now().UTC()
|
||||
from, err := queryDate(request, "from", now.AddDate(0, 0, -29))
|
||||
if err != nil {
|
||||
apiresponse.Error(writer, http.StatusBadRequest, "from 日期无效")
|
||||
return
|
||||
}
|
||||
to, err := queryDate(request, "to", now)
|
||||
if err != nil || to.Before(from) || to.Sub(from) > 366*24*time.Hour {
|
||||
apiresponse.Error(writer, http.StatusBadRequest, "usage 查询日期范围无效或超过 366 天")
|
||||
return
|
||||
}
|
||||
items, err := h.query.ListDailyUsage(request.Context(), UsageFilter{
|
||||
From: from, To: to, APIKeyID: strings.TrimSpace(request.URL.Query().Get("api_key_id")),
|
||||
Provider: strings.TrimSpace(request.URL.Query().Get("provider")), Model: strings.TrimSpace(request.URL.Query().Get("model")),
|
||||
})
|
||||
if err != nil {
|
||||
apiresponse.Error(writer, http.StatusServiceUnavailable, "usage 查询服务暂不可用")
|
||||
return
|
||||
}
|
||||
apiresponse.OK(writer, items)
|
||||
}
|
||||
|
||||
func (h *AdminHTTPHandler) requirePermission(writer http.ResponseWriter, request *http.Request, permission string) bool {
|
||||
account, err := h.identity.Authenticate(request.Context(), identity.KindAdmin, request.Header.Get("Authorization"))
|
||||
if err != nil {
|
||||
status := http.StatusUnauthorized
|
||||
if !errors.Is(err, identity.ErrInvalidSession) && !errors.Is(err, identity.ErrNotFound) {
|
||||
status = http.StatusServiceUnavailable
|
||||
}
|
||||
apiresponse.Error(writer, status, "登录状态无效或身份服务暂不可用")
|
||||
return false
|
||||
}
|
||||
if !identity.HasPermission(account, permission) {
|
||||
apiresponse.Error(writer, http.StatusForbidden, "缺少审计或 usage 查看权限")
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func queryTime(request *http.Request, name string, fallback time.Time) (time.Time, error) {
|
||||
value := request.URL.Query().Get(name)
|
||||
if value == "" {
|
||||
return fallback, nil
|
||||
}
|
||||
return time.Parse(time.RFC3339, value)
|
||||
}
|
||||
|
||||
func queryDate(request *http.Request, name string, fallback time.Time) (time.Time, error) {
|
||||
value := request.URL.Query().Get(name)
|
||||
if value == "" {
|
||||
value = fallback.Format("2006-01-02")
|
||||
}
|
||||
return time.Parse("2006-01-02", value)
|
||||
}
|
||||
@@ -0,0 +1,147 @@
|
||||
package audit
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"sort"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
const maintenanceLockID int64 = 6720240816
|
||||
|
||||
var auditPartitionPattern = regexp.MustCompile(`^audit_events_(\d{4})(\d{2})$`)
|
||||
|
||||
type MaintenanceResult struct {
|
||||
CreatedPartitions []string `json:"created_partitions"`
|
||||
DroppedPartitions []string `json:"dropped_partitions"`
|
||||
DeletedAuditRows int64 `json:"deleted_audit_rows"`
|
||||
DeletedUsageRows int64 `json:"deleted_usage_rows"`
|
||||
}
|
||||
|
||||
type Maintenance struct {
|
||||
pool *pgxpool.Pool
|
||||
auditRetention time.Duration
|
||||
usageRetention time.Duration
|
||||
monthsAhead int
|
||||
}
|
||||
|
||||
func NewMaintenance(pool *pgxpool.Pool, auditRetention, usageRetention time.Duration, monthsAhead int) *Maintenance {
|
||||
return &Maintenance{pool: pool, auditRetention: auditRetention, usageRetention: usageRetention, monthsAhead: monthsAhead}
|
||||
}
|
||||
|
||||
func (m *Maintenance) Run(ctx context.Context, now time.Time) (MaintenanceResult, error) {
|
||||
var result MaintenanceResult
|
||||
if m == nil || m.pool == nil {
|
||||
return result, errors.New("audit maintenance store unavailable")
|
||||
}
|
||||
now = now.UTC()
|
||||
auditCutoff := now.Add(-m.auditRetention)
|
||||
usageCutoff := now.Add(-m.usageRetention)
|
||||
tx, err := m.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return result, fmt.Errorf("begin audit maintenance: %w", err)
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock($1)`, maintenanceLockID); err != nil {
|
||||
return result, fmt.Errorf("lock audit maintenance: %w", err)
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `LOCK TABLE gateway.audit_events IN ACCESS EXCLUSIVE MODE`); err != nil {
|
||||
return result, fmt.Errorf("lock audit table: %w", err)
|
||||
}
|
||||
partitions, err := listAuditPartitions(ctx, tx)
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
for name := range partitions {
|
||||
start, ok := auditPartitionMonth(name)
|
||||
if !ok || start.AddDate(0, 1, 0).After(auditCutoff) {
|
||||
continue
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `DROP TABLE `+pgx.Identifier{"gateway", name}.Sanitize()); err != nil {
|
||||
return result, fmt.Errorf("drop audit partition %s: %w", name, err)
|
||||
}
|
||||
result.DroppedPartitions = append(result.DroppedPartitions, name)
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `ALTER TABLE gateway.audit_events DETACH PARTITION gateway.audit_events_default`); err != nil {
|
||||
return result, fmt.Errorf("detach audit default partition: %w", err)
|
||||
}
|
||||
deleted, err := tx.Exec(ctx, `DELETE FROM gateway.audit_events_default WHERE recorded_at < $1`, auditCutoff)
|
||||
if err != nil {
|
||||
return result, fmt.Errorf("clean audit default partition: %w", err)
|
||||
}
|
||||
result.DeletedAuditRows += deleted.RowsAffected()
|
||||
from := monthStart(auditCutoff)
|
||||
through := monthStart(now).AddDate(0, m.monthsAhead+1, 0)
|
||||
for start := from; start.Before(through); start = start.AddDate(0, 1, 0) {
|
||||
end := start.AddDate(0, 1, 0)
|
||||
name := "audit_events_" + start.Format("200601")
|
||||
if _, exists := partitions[name]; !exists {
|
||||
statement := fmt.Sprintf(`CREATE TABLE %s PARTITION OF gateway.audit_events FOR VALUES FROM ('%s') TO ('%s')`, pgx.Identifier{"gateway", name}.Sanitize(), start.Format(time.RFC3339), end.Format(time.RFC3339))
|
||||
if _, err := tx.Exec(ctx, statement); err != nil {
|
||||
return result, fmt.Errorf("create audit partition %s: %w", name, err)
|
||||
}
|
||||
result.CreatedPartitions = append(result.CreatedPartitions, name)
|
||||
}
|
||||
statement := `WITH moved AS (DELETE FROM gateway.audit_events_default WHERE recorded_at >= $1 AND recorded_at < $2 RETURNING *) INSERT INTO gateway.audit_events SELECT * FROM moved`
|
||||
if _, err := tx.Exec(ctx, statement, start, end); err != nil {
|
||||
return result, fmt.Errorf("move default audit rows into %s: %w", name, err)
|
||||
}
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `ALTER TABLE gateway.audit_events ATTACH PARTITION gateway.audit_events_default DEFAULT`); err != nil {
|
||||
return result, fmt.Errorf("reattach audit default partition: %w", err)
|
||||
}
|
||||
deleted, err = tx.Exec(ctx, `DELETE FROM gateway.audit_events WHERE recorded_at < $1`, auditCutoff)
|
||||
if err != nil {
|
||||
return result, fmt.Errorf("apply exact audit retention: %w", err)
|
||||
}
|
||||
result.DeletedAuditRows += deleted.RowsAffected()
|
||||
deleted, err = tx.Exec(ctx, `DELETE FROM gateway.usage_daily WHERE usage_date < $1::date`, usageCutoff.Format("2006-01-02"))
|
||||
if err != nil {
|
||||
return result, fmt.Errorf("apply usage retention: %w", err)
|
||||
}
|
||||
result.DeletedUsageRows = deleted.RowsAffected()
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return result, fmt.Errorf("commit audit maintenance: %w", err)
|
||||
}
|
||||
sort.Strings(result.CreatedPartitions)
|
||||
sort.Strings(result.DroppedPartitions)
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func listAuditPartitions(ctx context.Context, tx pgx.Tx) (map[string]struct{}, error) {
|
||||
rows, err := tx.Query(ctx, `SELECT child.relname FROM pg_inherits i JOIN pg_class parent ON parent.oid=i.inhparent JOIN pg_namespace n ON n.oid=parent.relnamespace JOIN pg_class child ON child.oid=i.inhrelid WHERE n.nspname='gateway' AND parent.relname='audit_events'`)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list audit partitions: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
result := make(map[string]struct{})
|
||||
for rows.Next() {
|
||||
var name string
|
||||
if err := rows.Scan(&name); err != nil {
|
||||
return nil, fmt.Errorf("scan audit partition: %w", err)
|
||||
}
|
||||
if name != "audit_events_default" {
|
||||
result[name] = struct{}{}
|
||||
}
|
||||
}
|
||||
return result, rows.Err()
|
||||
}
|
||||
|
||||
func auditPartitionMonth(name string) (time.Time, bool) {
|
||||
match := auditPartitionPattern.FindStringSubmatch(name)
|
||||
if match == nil {
|
||||
return time.Time{}, false
|
||||
}
|
||||
parsed, err := time.Parse("200601", match[1]+match[2])
|
||||
return parsed.UTC(), err == nil
|
||||
}
|
||||
|
||||
func monthStart(value time.Time) time.Time {
|
||||
value = value.UTC()
|
||||
return time.Date(value.Year(), value.Month(), 1, 0, 0, 0, 0, time.UTC)
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
package audit
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
func TestMaintenancePartitionsRetentionAndIdempotency(t *testing.T) {
|
||||
databaseURL := os.Getenv("AUDIT_MAINTENANCE_TEST_DATABASE_URL")
|
||||
if databaseURL == "" {
|
||||
t.Skip("AUDIT_MAINTENANCE_TEST_DATABASE_URL is not configured")
|
||||
}
|
||||
ctx := context.Background()
|
||||
pool, err := pgxpool.New(ctx, databaseURL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer pool.Close()
|
||||
const oldID = "44444444-4444-4444-8444-444444444444"
|
||||
const currentID = "55555555-5555-4555-8555-555555555555"
|
||||
_, err = pool.Exec(ctx, `INSERT INTO gateway.audit_events(id,request_id,protocol,status_code,recorded_at) VALUES
|
||||
($1,'maintenance-old','/v1/test',200,'2026-01-15T00:00:00Z'),
|
||||
($2,'maintenance-current','/v1/test',200,'2026-08-10T00:00:00Z')`, oldID, currentID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
now := time.Date(2026, 8, 11, 0, 0, 0, 0, time.UTC)
|
||||
first, err := NewMaintenance(pool, 300*24*time.Hour, 730*24*time.Hour, 1).Run(ctx, now)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !contains(first.CreatedPartitions, "audit_events_202601") || !contains(first.CreatedPartitions, "audit_events_202608") {
|
||||
t.Fatalf("expected old and current partitions, got %#v", first.CreatedPartitions)
|
||||
}
|
||||
second, err := NewMaintenance(pool, 90*24*time.Hour, 730*24*time.Hour, 1).Run(ctx, now)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !contains(second.DroppedPartitions, "audit_events_202601") {
|
||||
t.Fatalf("expected stale partition to be dropped, got %#v", second.DroppedPartitions)
|
||||
}
|
||||
var currentTable string
|
||||
if err := pool.QueryRow(ctx, `SELECT tableoid::regclass::text FROM gateway.audit_events WHERE id=$1`, currentID).Scan(¤tTable); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if currentTable != "gateway.audit_events_202608" && currentTable != "audit_events_202608" {
|
||||
t.Fatalf("current row was not routed to monthly partition: %s", currentTable)
|
||||
}
|
||||
var oldCount int
|
||||
if err := pool.QueryRow(ctx, `SELECT count(*) FROM gateway.audit_events WHERE id=$1`, oldID).Scan(&oldCount); err != nil || oldCount != 0 {
|
||||
t.Fatalf("expired row still exists: count=%d err=%v", oldCount, err)
|
||||
}
|
||||
third, err := NewMaintenance(pool, 90*24*time.Hour, 730*24*time.Hour, 1).Run(ctx, now)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(third.CreatedPartitions) != 0 || len(third.DroppedPartitions) != 0 {
|
||||
t.Fatalf("maintenance must be idempotent: %#v", third)
|
||||
}
|
||||
}
|
||||
|
||||
func contains(values []string, wanted string) bool {
|
||||
for _, value := range values {
|
||||
if value == wanted {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
package audit
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestAuditPartitionMonth(t *testing.T) {
|
||||
month, ok := auditPartitionMonth("audit_events_202608")
|
||||
if !ok || month.Format("2006-01-02") != "2026-08-01" {
|
||||
t.Fatalf("unexpected month: %v %v", month, ok)
|
||||
}
|
||||
if _, ok := auditPartitionMonth("audit_events_default"); ok {
|
||||
t.Fatal("default partition must not parse as a monthly partition")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMonthStartUsesUTC(t *testing.T) {
|
||||
value := time.Date(2026, 8, 31, 23, 0, 0, 0, time.FixedZone("UTC-2", -2*3600))
|
||||
if got := monthStart(value); got.Format(time.RFC3339) != "2026-09-01T00:00:00Z" {
|
||||
t.Fatalf("unexpected UTC month start: %s", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,151 @@
|
||||
package audit
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
type QueryService struct{ pool *pgxpool.Pool }
|
||||
|
||||
func NewQueryService(pool *pgxpool.Pool) *QueryService { return &QueryService{pool: pool} }
|
||||
|
||||
type EventView struct {
|
||||
ID string `json:"id"`
|
||||
TenantID *string `json:"tenant_id"`
|
||||
RequestID string `json:"request_id"`
|
||||
APIKeyID *string `json:"api_key_id"`
|
||||
APIKeyName *string `json:"api_key_name"`
|
||||
ProviderCode *string `json:"provider_code"`
|
||||
Model *string `json:"model"`
|
||||
Protocol string `json:"protocol"`
|
||||
StatusCode *int `json:"status_code"`
|
||||
PromptTokens *int64 `json:"prompt_tokens"`
|
||||
CompletionTokens *int64 `json:"completion_tokens"`
|
||||
CostMicrounits *int64 `json:"cost_microunits"`
|
||||
LatencyMS *int `json:"latency_ms"`
|
||||
Labels json.RawMessage `json:"labels"`
|
||||
RecordedAt time.Time `json:"recorded_at"`
|
||||
}
|
||||
|
||||
type EventFilter struct {
|
||||
From, To time.Time
|
||||
Before *time.Time
|
||||
APIKeyID string
|
||||
Provider string
|
||||
Model string
|
||||
StatusCode *int
|
||||
Limit int
|
||||
}
|
||||
|
||||
func (s *QueryService) ListEvents(ctx context.Context, filter EventFilter) ([]EventView, error) {
|
||||
if s == nil || s.pool == nil {
|
||||
return nil, fmt.Errorf("audit store unavailable")
|
||||
}
|
||||
where := []string{"a.recorded_at >= $1", "a.recorded_at < $2"}
|
||||
args := []any{filter.From, filter.To}
|
||||
add := func(condition string, value any) {
|
||||
args = append(args, value)
|
||||
where = append(where, fmt.Sprintf(condition, len(args)))
|
||||
}
|
||||
if filter.Before != nil {
|
||||
add("a.recorded_at < $%d", *filter.Before)
|
||||
}
|
||||
if filter.APIKeyID != "" {
|
||||
add("a.api_key_id = $%d", filter.APIKeyID)
|
||||
}
|
||||
if filter.Provider != "" {
|
||||
add("a.provider_code = $%d", filter.Provider)
|
||||
}
|
||||
if filter.Model != "" {
|
||||
add("a.model = $%d", filter.Model)
|
||||
}
|
||||
if filter.StatusCode != nil {
|
||||
add("a.status_code = $%d", *filter.StatusCode)
|
||||
}
|
||||
args = append(args, filter.Limit)
|
||||
query := `SELECT a.id::text,a.tenant_id::text,a.request_id,a.api_key_id::text,k.name,
|
||||
a.provider_code,a.model,a.protocol,a.status_code,a.prompt_tokens,a.completion_tokens,a.cost_microunits,a.latency_ms,a.labels,a.recorded_at
|
||||
FROM gateway.audit_events a LEFT JOIN gateway.api_keys k ON k.id=a.api_key_id
|
||||
WHERE ` + strings.Join(where, " AND ") + ` ORDER BY a.recorded_at DESC,a.id DESC LIMIT $` + strconv.Itoa(len(args))
|
||||
rows, err := s.pool.Query(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query audit events: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
items := make([]EventView, 0)
|
||||
for rows.Next() {
|
||||
var item EventView
|
||||
if err := rows.Scan(&item.ID, &item.TenantID, &item.RequestID, &item.APIKeyID, &item.APIKeyName,
|
||||
&item.ProviderCode, &item.Model, &item.Protocol, &item.StatusCode, &item.PromptTokens,
|
||||
&item.CompletionTokens, &item.CostMicrounits, &item.LatencyMS, &item.Labels, &item.RecordedAt); err != nil {
|
||||
return nil, fmt.Errorf("scan audit event: %w", err)
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
type DailyUsageView struct {
|
||||
Date time.Time `json:"date"`
|
||||
APIKeyID string `json:"api_key_id"`
|
||||
APIKeyName string `json:"api_key_name"`
|
||||
ProviderCode string `json:"provider_code"`
|
||||
Model string `json:"model"`
|
||||
Requests int64 `json:"requests"`
|
||||
FailedRequests int64 `json:"failed_requests"`
|
||||
PromptTokens int64 `json:"prompt_tokens"`
|
||||
CompletionTokens int64 `json:"completion_tokens"`
|
||||
CostMicrounits int64 `json:"cost_microunits"`
|
||||
}
|
||||
|
||||
type UsageFilter struct {
|
||||
From, To time.Time
|
||||
APIKeyID string
|
||||
Provider string
|
||||
Model string
|
||||
}
|
||||
|
||||
func (s *QueryService) ListDailyUsage(ctx context.Context, filter UsageFilter) ([]DailyUsageView, error) {
|
||||
if s == nil || s.pool == nil {
|
||||
return nil, fmt.Errorf("usage store unavailable")
|
||||
}
|
||||
where := []string{"u.usage_date >= $1", "u.usage_date <= $2"}
|
||||
args := []any{filter.From, filter.To}
|
||||
add := func(column string, value any) {
|
||||
args = append(args, value)
|
||||
where = append(where, column+" = $"+strconv.Itoa(len(args)))
|
||||
}
|
||||
if filter.APIKeyID != "" {
|
||||
add("u.api_key_id", filter.APIKeyID)
|
||||
}
|
||||
if filter.Provider != "" {
|
||||
add("u.provider_code", filter.Provider)
|
||||
}
|
||||
if filter.Model != "" {
|
||||
add("u.model", filter.Model)
|
||||
}
|
||||
rows, err := s.pool.Query(ctx, `SELECT u.usage_date,u.api_key_id::text,k.name,u.provider_code,u.model,
|
||||
u.requests,u.failed_requests,u.prompt_tokens,u.completion_tokens,u.cost_microunits
|
||||
FROM gateway.usage_daily u JOIN gateway.api_keys k ON k.id=u.api_key_id WHERE `+
|
||||
strings.Join(where, " AND ")+` ORDER BY u.usage_date DESC,k.name,u.provider_code,u.model`, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query daily usage: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
items := make([]DailyUsageView, 0)
|
||||
for rows.Next() {
|
||||
var item DailyUsageView
|
||||
if err := rows.Scan(&item.Date, &item.APIKeyID, &item.APIKeyName, &item.ProviderCode, &item.Model,
|
||||
&item.Requests, &item.FailedRequests, &item.PromptTokens, &item.CompletionTokens, &item.CostMicrounits); err != nil {
|
||||
return nil, fmt.Errorf("scan daily usage: %w", err)
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
@@ -0,0 +1,282 @@
|
||||
package audit
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
platformid "aigateway.local/core/internal/platform/id"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgtype"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
type Event struct {
|
||||
TenantID *string
|
||||
RequestID string
|
||||
APIKeyID *string
|
||||
ProviderCode string
|
||||
Model string
|
||||
Protocol string
|
||||
StatusCode int
|
||||
PromptTokens int64
|
||||
CompletionTokens int64
|
||||
CostMicrounits int64
|
||||
PriceID string
|
||||
Currency string
|
||||
LatencyMS int
|
||||
Labels map[string]any
|
||||
RecordedAt time.Time
|
||||
}
|
||||
|
||||
type Stats struct {
|
||||
Accepted uint64
|
||||
Dropped uint64
|
||||
Written uint64
|
||||
Failures uint64
|
||||
}
|
||||
|
||||
type Recorder struct {
|
||||
pool *pgxpool.Pool
|
||||
logger *slog.Logger
|
||||
queue chan Event
|
||||
batchSize int
|
||||
flushInterval time.Duration
|
||||
accepted atomic.Uint64
|
||||
dropped atomic.Uint64
|
||||
written atomic.Uint64
|
||||
failures atomic.Uint64
|
||||
}
|
||||
|
||||
func NewRecorder(pool *pgxpool.Pool, logger *slog.Logger, queueSize, batchSize int, flushInterval time.Duration) *Recorder {
|
||||
if queueSize < 1 {
|
||||
queueSize = 4096
|
||||
}
|
||||
if batchSize < 1 {
|
||||
batchSize = 200
|
||||
}
|
||||
if flushInterval <= 0 {
|
||||
flushInterval = time.Second
|
||||
}
|
||||
return &Recorder{pool: pool, logger: logger, queue: make(chan Event, queueSize), batchSize: batchSize, flushInterval: flushInterval}
|
||||
}
|
||||
|
||||
func (r *Recorder) Record(event Event) bool {
|
||||
if r == nil || r.pool == nil {
|
||||
return false
|
||||
}
|
||||
if event.RecordedAt.IsZero() {
|
||||
event.RecordedAt = time.Now().UTC()
|
||||
}
|
||||
select {
|
||||
case r.queue <- event:
|
||||
r.accepted.Add(1)
|
||||
return true
|
||||
default:
|
||||
r.dropped.Add(1)
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Recorder) Stats() Stats {
|
||||
if r == nil {
|
||||
return Stats{}
|
||||
}
|
||||
return Stats{Accepted: r.accepted.Load(), Dropped: r.dropped.Load(), Written: r.written.Load(), Failures: r.failures.Load()}
|
||||
}
|
||||
|
||||
func (r *Recorder) Prometheus() string {
|
||||
stats := r.Stats()
|
||||
return fmt.Sprintf("gateway_audit_events_accepted_total %d\ngateway_audit_events_dropped_total %d\ngateway_audit_events_written_total %d\ngateway_audit_flush_failures_total %d\n", stats.Accepted, stats.Dropped, stats.Written, stats.Failures)
|
||||
}
|
||||
|
||||
func (r *Recorder) Run(ctx context.Context) {
|
||||
ticker := time.NewTicker(r.flushInterval)
|
||||
defer ticker.Stop()
|
||||
pending := make([]Event, 0, r.batchSize)
|
||||
flushAllowed := true
|
||||
for {
|
||||
var events <-chan Event = r.queue
|
||||
if len(pending) >= cap(r.queue)+r.batchSize {
|
||||
events = nil
|
||||
}
|
||||
select {
|
||||
case event := <-events:
|
||||
pending = append(pending, event)
|
||||
if len(pending) >= r.batchSize && flushAllowed {
|
||||
if r.flush(ctx, pending) == nil {
|
||||
pending = pending[:0]
|
||||
} else {
|
||||
flushAllowed = false
|
||||
}
|
||||
}
|
||||
case <-ticker.C:
|
||||
flushAllowed = true
|
||||
if len(pending) > 0 {
|
||||
if r.flush(ctx, pending) == nil {
|
||||
pending = pending[:0]
|
||||
} else {
|
||||
flushAllowed = false
|
||||
}
|
||||
}
|
||||
case <-ctx.Done():
|
||||
for {
|
||||
select {
|
||||
case event := <-r.queue:
|
||||
pending = append(pending, event)
|
||||
default:
|
||||
flushCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
if len(pending) > 0 {
|
||||
_ = r.flush(flushCtx, pending)
|
||||
}
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type dailyKey struct {
|
||||
date, apiKeyID, provider, model string
|
||||
}
|
||||
|
||||
type dailyValue struct {
|
||||
requests, failed, prompt, completion, cost int64
|
||||
}
|
||||
|
||||
type policyAlert struct {
|
||||
eventID, requestID string
|
||||
tenantID *string
|
||||
payload []byte
|
||||
}
|
||||
|
||||
func (r *Recorder) flush(ctx context.Context, events []Event) error {
|
||||
if len(events) == 0 {
|
||||
return nil
|
||||
}
|
||||
tx, err := r.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return r.flushError(err)
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
rows := make([][]any, 0, len(events))
|
||||
daily := make(map[dailyKey]dailyValue)
|
||||
alerts := make([]policyAlert, 0)
|
||||
for _, event := range events {
|
||||
id, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return r.flushError(err)
|
||||
}
|
||||
labels, _ := json.Marshal(event.Labels)
|
||||
if policies, matched := event.Labels["content_policies"]; matched {
|
||||
alertID, idErr := platformid.NewUUID()
|
||||
if idErr != nil {
|
||||
return r.flushError(idErr)
|
||||
}
|
||||
payload, _ := json.Marshal(map[string]any{
|
||||
"request_id": event.RequestID, "protocol": event.Protocol, "status_code": event.StatusCode,
|
||||
"provider_code": event.ProviderCode, "model": event.Model, "content_policies": policies,
|
||||
"content_redacted": event.Labels["content_redacted"], "recorded_at": event.RecordedAt,
|
||||
})
|
||||
alerts = append(alerts, policyAlert{eventID: alertID, requestID: event.RequestID, tenantID: event.TenantID, payload: payload})
|
||||
}
|
||||
rows = append(rows, []any{
|
||||
uuidValue(id), uuidPointer(event.TenantID), event.RequestID, nil, uuidPointer(event.APIKeyID),
|
||||
nullString(event.ProviderCode), nullString(event.Model), event.Protocol, event.StatusCode,
|
||||
nullToken(event.PromptTokens), nullToken(event.CompletionTokens), nullCost(event.CostMicrounits), event.LatencyMS,
|
||||
nil, nil, labels, event.RecordedAt,
|
||||
})
|
||||
if event.APIKeyID != nil && *event.APIKeyID != "" {
|
||||
key := dailyKey{date: event.RecordedAt.UTC().Format("2006-01-02"), apiKeyID: *event.APIKeyID, provider: event.ProviderCode, model: event.Model}
|
||||
value := daily[key]
|
||||
value.requests++
|
||||
if event.StatusCode >= 400 {
|
||||
value.failed++
|
||||
}
|
||||
value.prompt += max(event.PromptTokens, 0)
|
||||
value.completion += max(event.CompletionTokens, 0)
|
||||
value.cost += max(event.CostMicrounits, 0)
|
||||
daily[key] = value
|
||||
}
|
||||
}
|
||||
_, err = tx.CopyFrom(ctx, pgx.Identifier{"gateway", "audit_events"}, []string{
|
||||
"id", "tenant_id", "request_id", "actor_id", "api_key_id", "provider_code", "model", "protocol",
|
||||
"status_code", "prompt_tokens", "completion_tokens", "cost_microunits", "latency_ms",
|
||||
"request_preview", "response_preview", "labels", "recorded_at",
|
||||
}, pgx.CopyFromRows(rows))
|
||||
if err != nil {
|
||||
return r.flushError(err)
|
||||
}
|
||||
batch := &pgx.Batch{}
|
||||
for key, value := range daily {
|
||||
batch.Queue(`INSERT INTO gateway.usage_daily(usage_date,api_key_id,provider_code,model,requests,failed_requests,prompt_tokens,completion_tokens,cost_microunits)
|
||||
VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9)
|
||||
ON CONFLICT(usage_date,api_key_id,provider_code,model) DO UPDATE SET
|
||||
requests=gateway.usage_daily.requests+EXCLUDED.requests,
|
||||
failed_requests=gateway.usage_daily.failed_requests+EXCLUDED.failed_requests,
|
||||
prompt_tokens=gateway.usage_daily.prompt_tokens+EXCLUDED.prompt_tokens,
|
||||
completion_tokens=gateway.usage_daily.completion_tokens+EXCLUDED.completion_tokens,
|
||||
cost_microunits=gateway.usage_daily.cost_microunits+EXCLUDED.cost_microunits,
|
||||
updated_at=clock_timestamp()`, key.date, key.apiKeyID, key.provider, key.model, value.requests, value.failed, value.prompt, value.completion, value.cost)
|
||||
}
|
||||
for _, alert := range alerts {
|
||||
batch.Queue(`INSERT INTO gateway.outbox_events(event_id,event_type,event_version,tenant_id,aggregate_type,aggregate_id,payload)
|
||||
VALUES($1,'content_policy.matched',1,$2,'gateway_request',$3,$4)`, alert.eventID, uuidPointer(alert.tenantID), alert.requestID, alert.payload)
|
||||
}
|
||||
results := tx.SendBatch(ctx, batch)
|
||||
if err := results.Close(); err != nil {
|
||||
return r.flushError(err)
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return r.flushError(err)
|
||||
}
|
||||
r.written.Add(uint64(len(events)))
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *Recorder) flushError(err error) error {
|
||||
r.failures.Add(1)
|
||||
if r.logger != nil {
|
||||
r.logger.Error("audit batch flush failed", "error", err)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func uuidValue(value string) pgtype.UUID {
|
||||
var result pgtype.UUID
|
||||
_ = result.Scan(value)
|
||||
return result
|
||||
}
|
||||
|
||||
func uuidPointer(value *string) any {
|
||||
if value == nil || strings.TrimSpace(*value) == "" {
|
||||
return nil
|
||||
}
|
||||
return uuidValue(*value)
|
||||
}
|
||||
|
||||
func nullString(value string) any {
|
||||
if value == "" {
|
||||
return nil
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func nullToken(value int64) any {
|
||||
if value <= 0 {
|
||||
return nil
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func nullCost(value int64) any {
|
||||
if value <= 0 {
|
||||
return nil
|
||||
}
|
||||
return value
|
||||
}
|
||||
Reference in New Issue
Block a user