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:
ben
2026-08-12 11:45:54 +08:00
commit 5759c1862e
807 changed files with 114727 additions and 0 deletions
+383
View File
@@ -0,0 +1,383 @@
package main
import (
"context"
"errors"
"log/slog"
"net/http"
"os"
"os/signal"
"syscall"
"time"
"aigateway.local/core/internal/apikey"
"aigateway.local/core/internal/audit"
"aigateway.local/core/internal/contentpolicy"
"aigateway.local/core/internal/factcheck"
"aigateway.local/core/internal/gateway"
"aigateway.local/core/internal/identity"
"aigateway.local/core/internal/operations"
"aigateway.local/core/internal/outbox"
"aigateway.local/core/internal/platform/cache"
"aigateway.local/core/internal/platform/config"
"aigateway.local/core/internal/platform/cryptox"
"aigateway.local/core/internal/platform/database"
"aigateway.local/core/internal/platform/health"
"aigateway.local/core/internal/platform/httpserver"
"aigateway.local/core/internal/portal"
"aigateway.local/core/internal/pricing"
"aigateway.local/core/internal/provider"
providercontrolplane "aigateway.local/core/internal/provider/controlplane"
provideropenai "aigateway.local/core/internal/provider/openai"
providerruntime "aigateway.local/core/internal/provider/runtime"
"aigateway.local/core/internal/shadow"
"aigateway.local/core/internal/workbench"
)
var version = "dev"
func main() {
logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo}))
cfg, err := config.Load()
if err != nil {
logger.Error("invalid configuration", "error", err)
os.Exit(1)
}
if err := cfg.ValidateRuntime(); err != nil {
logger.Error("invalid runtime configuration", "error", err)
os.Exit(1)
}
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
defer stop()
db, err := database.Open(ctx, cfg.Database)
if err != nil {
logger.Error("database initialization failed", "error", err)
os.Exit(1)
}
if db != nil {
defer db.Close()
}
criticalRedis, err := cache.Open(cfg.Redis.CriticalURL)
if err != nil {
logger.Error("critical redis initialization failed", "error", err)
os.Exit(1)
}
if criticalRedis != nil {
defer criticalRedis.Close()
}
cacheRedis, err := cache.Open(cfg.Redis.CacheURL)
if err != nil {
logger.Error("cache redis initialization failed", "error", err)
os.Exit(1)
}
if cacheRedis != nil {
defer cacheRedis.Close()
}
checker := health.Checker{Timeout: 1500 * time.Millisecond}
checker.Dependencies = append(checker.Dependencies,
health.Dependency{Name: "postgres", Required: true, Probe: func(probeCtx context.Context) error {
if db == nil {
return errors.New("not configured")
}
return db.Ping(probeCtx)
}},
health.Dependency{Name: "redis_critical", Required: true, Probe: func(probeCtx context.Context) error {
if criticalRedis == nil {
return errors.New("not configured")
}
return criticalRedis.Ping(probeCtx).Err()
}},
health.Dependency{Name: "redis_cache", Required: false, Probe: func(probeCtx context.Context) error {
if cacheRedis == nil {
return errors.New("not configured")
}
return cacheRedis.Ping(probeCtx).Err()
}},
)
adapter, err := provideropenai.New(cfg.Upstream.BaseURL, cfg.Upstream.APIKey)
if err != nil {
logger.Error("provider initialization failed", "error", err)
os.Exit(1)
}
var fallbackAdapter provider.Adapter
if cfg.Upstream.FallbackEnabled {
fallbackAdapter = adapter
}
credentialCipher, err := provider.NewCredentialCipher(
cfg.Credentials.MasterKey, cfg.Credentials.KEKVersion, cfg.Credentials.KEKKeyring,
)
if err != nil {
logger.Error("credential encryption initialization failed", "error", err)
os.Exit(1)
}
providerRepository := provider.NewRepository(db)
providerResolver := providerruntime.NewResolver(
providerRepository, credentialCipher, fallbackAdapter, cfg.Credentials.ProviderRefreshInterval, logger,
)
providerResolver.SetNotificationClient(criticalRedis)
go providerResolver.Run(ctx)
apiKeyRepository := apikey.NewRepository(db)
bootstrapAPIKey := ""
if cfg.Security.BootstrapAPIKeyEnabled {
bootstrapAPIKey = cfg.Security.BootstrapAPIKey
logger.Warn("bootstrap API key compatibility is enabled; monitor usage and disable after migration")
}
apiKeyAuthenticator := apikey.NewAuthenticator(apiKeyRepository, criticalRedis, bootstrapAPIKey)
apiKeyAuthenticator.SetLogger(logger)
proxy := gateway.NewDynamicProxy(providerResolver, apiKeyAuthenticator, cfg.Server.MaxBodyBytes, logger)
proxy.SetAdmissionController(gateway.NewRedisAdmissionController(criticalRedis))
proxy.SetTokenQuotaController(gateway.NewRedisTokenQuotaController(criticalRedis))
proxy.SetResiliencePolicy(gateway.ResiliencePolicy{
ResponseHeaderTimeout: cfg.Upstream.ResponseHeaderTimeout, MaxRetries: cfg.Upstream.MaxRetries,
RetryBackoff: cfg.Upstream.RetryBackoff, CircuitThreshold: cfg.Upstream.CircuitThreshold,
CircuitOpenDuration: cfg.Upstream.CircuitOpenDuration,
})
auditRecorder := audit.NewRecorder(db, logger, cfg.Audit.QueueSize, cfg.Audit.BatchSize, cfg.Audit.FlushInterval)
auditContext, stopAudit := context.WithCancel(context.Background())
auditStopped := make(chan struct{})
go func() {
auditRecorder.Run(auditContext)
close(auditStopped)
}()
proxy.SetAuditRecorder(auditRecorder)
contentPolicyEngine := contentpolicy.NewEngine(db, cfg.RuntimeData.ContentPolicyRefreshInterval, logger)
pricingService := pricing.NewService(db, cfg.RuntimeData.PricingRefreshInterval, logger)
if err := contentPolicyEngine.Reload(ctx); err != nil {
logger.Error("content policy initialization failed", "error", err)
os.Exit(1)
}
if err := pricingService.Reload(ctx); err != nil {
logger.Error("model pricing initialization failed", "error", err)
os.Exit(1)
}
go contentPolicyEngine.Run(ctx)
go pricingService.Run(ctx)
proxy.SetContentPolicyEngine(contentPolicyEngine)
proxy.SetPricingService(pricingService)
identityRepository := identity.NewRepository(db)
sessionStore := identity.NewSessionStore(criticalRedis, cfg.Auth.SessionTTL)
loginLimiter := identity.NewLoginLimiter(criticalRedis, cfg.Auth.LoginRateLimitMax, cfg.Auth.LoginRateLimitWindow)
totpCipher, err := cryptox.NewKeyring(
cfg.Credentials.MasterKey, cfg.Credentials.KEKVersion, cfg.Credentials.KEKKeyring, "totp-secret",
)
if err != nil {
logger.Error("TOTP encryption initialization failed", "error", err)
os.Exit(1)
}
identityService := identity.NewService(identityRepository, sessionStore, loginLimiter, cfg.Auth, totpCipher)
idpCipher, err := cryptox.NewKeyring(
cfg.Credentials.MasterKey, cfg.Credentials.KEKVersion, cfg.Credentials.KEKKeyring, "identity-provider-credentials",
)
if err != nil {
logger.Error("identity provider encryption initialization failed", "error", err)
os.Exit(1)
}
identityService.SetIdentityProviderCipher(idpCipher, cfg.Credentials.AllowPrivateProviderURL)
identityHandler := identity.NewHTTPHandler(identityService)
identityManagementHandler := identity.NewManagementHTTPHandler(identityService)
providerHandler := provider.NewAdminHTTPHandler(
providerRepository, credentialCipher, identityService, cfg.Credentials.AllowPrivateProviderURL,
)
providerOperations := providercontrolplane.NewService(
providerRepository, credentialCipher, cfg.Credentials.AllowPrivateProviderURL,
)
providerHandler.SetOperations(providerOperations)
providerHandler.SetChangeHook(func(changeCtx context.Context) error {
reloadErr := providerResolver.Reload(changeCtx)
notifyErr := providerResolver.Notify(changeCtx)
if reloadErr != nil || notifyErr != nil {
logger.Warn("provider change propagation was incomplete", "reload_error", reloadErr, "notify_error", notifyErr)
}
return errors.Join(reloadErr, notifyErr)
})
apiKeyHandler := apikey.NewAdminHTTPHandler(apiKeyRepository, apiKeyAuthenticator, identityService)
apiKeyHandler.SetUsageStore(apikey.NewUsageStore(criticalRedis))
auditHandler := audit.NewAdminHTTPHandler(audit.NewQueryService(db), identityService)
outboxHandler := outbox.NewAdminHTTPHandler(outbox.NewStore(db), identityService)
contentPolicyHandler := contentpolicy.NewAdminHTTPHandler(contentpolicy.NewStore(db), contentPolicyEngine, identityService)
pricingHandler := pricing.NewAdminHTTPHandler(pricingService, identityService)
factCheckHandler := factcheck.NewAdminHTTPHandler(factcheck.NewService(db), identityService)
workbenchService := workbench.NewService(db)
toolCipher, err := cryptox.NewKeyring(
cfg.Credentials.MasterKey, cfg.Credentials.KEKVersion, cfg.Credentials.KEKKeyring, "tool-request-headers",
)
if err != nil {
logger.Error("tool credential encryption initialization failed", "error", err)
os.Exit(1)
}
notificationCipher, err := cryptox.NewKeyring(
cfg.Credentials.MasterKey, cfg.Credentials.KEKVersion, cfg.Credentials.KEKKeyring, "notification-signing-secret",
)
if err != nil {
logger.Error("notification encryption initialization failed", "error", err)
os.Exit(1)
}
toolService := workbench.NewToolService(workbenchService, toolCipher, cfg.Credentials.AllowPrivateToolURL)
notificationService := workbench.NewNotificationService(workbenchService, notificationCipher, cfg.Credentials.AllowPrivateWebhookURL)
workbenchHandler := workbench.NewAdminHTTPHandler(workbenchService, toolService, notificationService, identityService)
mcpServerCipher, err := cryptox.NewKeyring(
cfg.Credentials.MasterKey, cfg.Credentials.KEKVersion, cfg.Credentials.KEKKeyring, "mcp-server-headers",
)
if err != nil {
logger.Error("MCP server credential encryption initialization failed", "error", err)
os.Exit(1)
}
mcpServerService := workbench.NewMCPServerService(workbenchService, mcpServerCipher, cfg.Credentials.AllowPrivateToolURL)
skillService := workbench.NewSkillService(workbenchService)
digitalEmployeeService := workbench.NewDigitalEmployeeService(workbenchService, skillService, toolService, mcpServerService)
marketplaceService := workbench.NewMarketplaceService(workbenchService, mcpServerService, skillService, digitalEmployeeService)
mcpClient := workbench.NewMCPClient(cfg.Credentials.AllowPrivateToolURL, 60*time.Second)
marketplaceHandler := workbench.NewMarketplaceAdminHTTPHandler(marketplaceService, mcpServerService, skillService, digitalEmployeeService, mcpClient, identityService)
applicationKeyCipher, err := cryptox.NewKeyring(
cfg.Credentials.MasterKey, cfg.Credentials.KEKVersion, cfg.Credentials.KEKKeyring, "application-runtime-key",
)
if err != nil {
logger.Error("application runtime credential initialization failed", "error", err)
os.Exit(1)
}
shadowMiddleware := shadow.New(cfg.Shadow, logger)
governedGateway := shadowMiddleware.Wrap(proxy)
workbenchRuntime := workbench.NewRuntimeHTTPHandler(workbenchService, toolService, workbench.NewPostgreSQLRetriever(workbenchService), apiKeyAuthenticator, governedGateway, workbench.MarketplaceDeps{
MCPServers: mcpServerService,
Skills: skillService,
Employees: digitalEmployeeService,
Market: marketplaceService,
MCPClient: mcpClient,
})
workbenchRuntime.SetLogger(logger)
// Wire the fact-check engine: the admin fact-check settings/policies UI now
// actually governs application answers instead of being inert configuration.
factCheckEngine := factcheck.NewEngine(db, workbench.NewFactCheckRetriever(workbench.NewPostgreSQLRetriever(workbenchService)), logger)
workbenchRuntime.SetFactCheckEngine(factCheckEngine)
portalService := portal.NewService(db, workbenchService, toolService, identityService)
portalService.SetApplicationRuntime(portal.NewRuntimeCredentials(db, apiKeyRepository, applicationKeyCipher), workbenchRuntime)
portalService.SetMarketplace(marketplaceService)
portalHandler := portal.NewHTTPHandler(portalService, identityService)
portalAdminHandler := portal.NewAdminHTTPHandler(portalService, identityService)
startedAt := time.Now()
operationsHandler := operations.NewAdminHTTPHandler(db, identityService, version, startedAt, func(reloadCtx context.Context) error {
return errors.Join(providerResolver.Reload(reloadCtx), contentPolicyEngine.Reload(reloadCtx), pricingService.Reload(reloadCtx))
})
controlMux := http.NewServeMux()
controlMux.Handle("/api/v1/admin/providers", providerHandler)
controlMux.Handle("/api/v1/admin/providers/", providerHandler)
controlMux.Handle("/api/v1/admin/model-routes", providerHandler)
controlMux.Handle("/api/v1/admin/model-routes/", providerHandler)
controlMux.Handle("/api/v1/admin/api-keys", apiKeyHandler)
controlMux.Handle("/api/v1/admin/api-keys/", apiKeyHandler)
controlMux.Handle("/api/v1/admin/audit-events", auditHandler)
controlMux.Handle("/api/v1/admin/usage/", auditHandler)
controlMux.Handle("/api/v1/admin/outbox-events", outboxHandler)
controlMux.Handle("/api/v1/admin/outbox-events/", outboxHandler)
controlMux.Handle("/api/v1/admin/content-policies", contentPolicyHandler)
controlMux.Handle("/api/v1/admin/content-policies/", contentPolicyHandler)
controlMux.Handle("/api/v1/admin/model-prices", pricingHandler)
controlMux.Handle("/api/v1/admin/model-prices/", pricingHandler)
controlMux.Handle("/api/v1/admin/fact-check/", factCheckHandler)
controlMux.Handle("/api/v1/admin/prompt-categories", workbenchHandler)
controlMux.Handle("/api/v1/admin/prompt-categories/", workbenchHandler)
controlMux.Handle("/api/v1/admin/prompts", workbenchHandler)
controlMux.Handle("/api/v1/admin/prompts/", workbenchHandler)
controlMux.Handle("/api/v1/admin/knowledge-bases", workbenchHandler)
controlMux.Handle("/api/v1/admin/knowledge-bases/", workbenchHandler)
controlMux.Handle("/api/v1/admin/tools", workbenchHandler)
controlMux.Handle("/api/v1/admin/tools/", workbenchHandler)
controlMux.Handle("/api/v1/admin/applications", workbenchHandler)
controlMux.Handle("/api/v1/admin/applications/", workbenchHandler)
controlMux.Handle("/api/v1/admin/marketplace-categories", marketplaceHandler)
controlMux.Handle("/api/v1/admin/marketplace-categories/", marketplaceHandler)
controlMux.Handle("/api/v1/admin/mcp-servers", marketplaceHandler)
controlMux.Handle("/api/v1/admin/mcp-servers/", marketplaceHandler)
controlMux.Handle("/api/v1/admin/skills", marketplaceHandler)
controlMux.Handle("/api/v1/admin/skills/", marketplaceHandler)
controlMux.Handle("/api/v1/admin/digital-employees", marketplaceHandler)
controlMux.Handle("/api/v1/admin/digital-employees/", marketplaceHandler)
controlMux.Handle("/api/v1/admin/marketplace/", marketplaceHandler)
controlMux.Handle("/api/v1/portal/marketplace", portalHandler)
controlMux.Handle("/api/v1/portal/marketplace/", portalHandler)
controlMux.Handle("/api/v1/admin/notification-channels", workbenchHandler)
controlMux.Handle("/api/v1/admin/notification-channels/", workbenchHandler)
controlMux.Handle("/api/v1/admin/notification-deliveries", workbenchHandler)
controlMux.Handle("/api/v1/admin/notification-deliveries/", workbenchHandler)
controlMux.Handle("/api/v1/admin/models", portalAdminHandler)
controlMux.Handle("/api/v1/admin/model-requests", portalAdminHandler)
controlMux.Handle("/api/v1/admin/model-requests/", portalAdminHandler)
controlMux.Handle("/api/v1/admin/system-info", operationsHandler)
controlMux.Handle("/api/v1/admin/monitoring/overview", operationsHandler)
controlMux.Handle("/api/v1/admin/reload", operationsHandler)
controlMux.Handle("/api/v1/admin/identities/", identityManagementHandler)
controlMux.Handle("/api/v1/admin/departments", identityManagementHandler)
controlMux.Handle("/api/v1/admin/departments/", identityManagementHandler)
controlMux.Handle("/api/v1/admin/identity-providers", identityManagementHandler)
controlMux.Handle("/api/v1/admin/identity-providers/", identityManagementHandler)
controlMux.Handle("/api/v1/admin/saml-providers", identityManagementHandler)
controlMux.Handle("/api/v1/admin/saml-providers/", identityManagementHandler)
controlMux.Handle("/api/v1/portal/applications", portalHandler)
controlMux.Handle("/api/v1/portal/apps/", portalHandler)
controlMux.Handle("/api/v1/portal/catalog", portalHandler)
controlMux.Handle("/api/v1/portal/cost", portalHandler)
controlMux.Handle("/api/v1/portal/docs-info", portalHandler)
controlMux.Handle("/api/v1/portal/knowledge", portalHandler)
controlMux.Handle("/api/v1/portal/logs", portalHandler)
controlMux.Handle("/api/v1/portal/logs/", portalHandler)
controlMux.Handle("/api/v1/portal/model-requests", portalHandler)
controlMux.Handle("/api/v1/portal/model-requests/", portalHandler)
controlMux.Handle("/api/v1/portal/password", portalHandler)
controlMux.Handle("/api/v1/portal/prompts", portalHandler)
controlMux.Handle("/api/v1/portal/prompts/", portalHandler)
controlMux.Handle("/api/v1/portal/stats", portalHandler)
controlMux.Handle("/api/v1/portal/tools", portalHandler)
controlMux.Handle("/api/v1/", identityHandler)
publicMux := http.NewServeMux()
publicMux.Handle("/v1/prompts", workbenchRuntime)
publicMux.Handle("/v1/prompts/", workbenchRuntime)
publicMux.Handle("/v1/knowledge/", workbenchRuntime)
publicMux.Handle("/v1/tools", workbenchRuntime)
publicMux.Handle("/v1/tools/", workbenchRuntime)
publicMux.Handle("/v1/applications/", workbenchRuntime)
publicMux.Handle("/v1/skills/", workbenchRuntime)
publicMux.Handle("/v1/mcp-servers", workbenchRuntime)
publicMux.Handle("/v1/mcp-servers/", workbenchRuntime)
publicMux.Handle("/v1/digital-employees/", workbenchRuntime)
publicMux.Handle("/v1/", governedGateway)
server := httpserver.New(httpserver.Dependencies{
Config: cfg, Logger: logger, Checker: checker, Gateway: publicMux, Control: controlMux,
Version: version, StartedAt: startedAt,
BootstrapUses: apiKeyAuthenticator.BootstrapUses,
ExtraMetrics: func() string { return auditRecorder.Prometheus() + shadowMiddleware.Prometheus() },
})
serverErrors := make(chan error, 1)
go func() {
logger.Info("gateway API started", "address", server.Addr, "version", version)
serverErrors <- server.ListenAndServe()
}()
select {
case <-ctx.Done():
logger.Info("shutdown requested")
case err := <-serverErrors:
if !errors.Is(err, http.ErrServerClosed) {
logger.Error("gateway API stopped unexpectedly", "error", err)
os.Exit(1)
}
}
shutdownCtx, cancel := context.WithTimeout(context.Background(), cfg.Server.ShutdownTimeout)
defer cancel()
if err := server.Shutdown(shutdownCtx); err != nil {
logger.Error("graceful shutdown failed", "error", err)
_ = server.Close()
os.Exit(1)
}
stopAudit()
select {
case <-auditStopped:
case <-shutdownCtx.Done():
logger.Warn("audit recorder drain timed out")
}
logger.Info("gateway API stopped")
}
+84
View File
@@ -0,0 +1,84 @@
package main
import (
"context"
"flag"
"log/slog"
"os"
"strings"
"aigateway.local/core/internal/identity"
"aigateway.local/core/internal/platform/config"
"aigateway.local/core/internal/platform/database"
)
func main() {
adminUsername := flag.String("admin-username", env("BOOTSTRAP_ADMIN_USERNAME", "admin"), "initial administrator username")
adminDisplayName := flag.String("admin-display-name", "系统管理员", "initial administrator display name")
portalAccount := flag.String("portal-account", strings.TrimSpace(os.Getenv("BOOTSTRAP_PORTAL_ACCOUNT")), "optional initial portal account")
portalName := flag.String("portal-name", "初始用户", "initial portal user display name")
flag.Parse()
logger := slog.New(slog.NewJSONHandler(os.Stdout, nil))
cfg, err := config.Load()
if err != nil {
logger.Error("invalid configuration", "error", err)
os.Exit(1)
}
if cfg.Database.URL == "" {
logger.Error("DATABASE_URL is required")
os.Exit(1)
}
adminPassword := os.Getenv("BOOTSTRAP_ADMIN_PASSWORD")
portalPassword := os.Getenv("BOOTSTRAP_PORTAL_PASSWORD")
if len(adminPassword) < 12 {
logger.Error("BOOTSTRAP_ADMIN_PASSWORD must contain at least 12 characters")
os.Exit(1)
}
if *portalAccount != "" && len(portalPassword) < 12 {
logger.Error("BOOTSTRAP_PORTAL_PASSWORD must contain at least 12 characters when creating a portal user")
os.Exit(1)
}
ctx := context.Background()
pool, err := database.Open(ctx, cfg.Database)
if err != nil {
logger.Error("database initialization failed", "error", err)
os.Exit(1)
}
defer pool.Close()
repository := identity.NewRepository(pool)
hasher := identity.PasswordHasher{}
adminHash, err := hasher.Hash(adminPassword)
if err != nil {
logger.Error("password hashing failed", "error", err)
os.Exit(1)
}
adminID, err := repository.CreateAdmin(ctx, *adminUsername, *adminDisplayName, "superadmin", adminHash)
if err != nil {
logger.Error("administrator bootstrap failed", "error", err)
os.Exit(1)
}
logger.Info("administrator created", "id", adminID, "username", strings.ToLower(strings.TrimSpace(*adminUsername)))
if *portalAccount != "" {
portalHash, err := hasher.Hash(portalPassword)
if err != nil {
logger.Error("portal password hashing failed", "error", err)
os.Exit(1)
}
portalID, err := repository.CreatePortalUser(ctx, *portalAccount, *portalName, portalHash)
if err != nil {
logger.Error("portal user bootstrap failed", "error", err)
os.Exit(1)
}
logger.Info("portal user created", "id", portalID, "account", strings.ToLower(strings.TrimSpace(*portalAccount)))
}
}
func env(key, fallback string) string {
if value := strings.TrimSpace(os.Getenv(key)); value != "" {
return value
}
return fallback
}
+178
View File
@@ -0,0 +1,178 @@
package main
import (
"bufio"
"context"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"errors"
"flag"
"fmt"
"io"
"os"
"strings"
"aigateway.local/core/internal/platform/config"
"aigateway.local/core/internal/platform/database"
platformid "aigateway.local/core/internal/platform/id"
"aigateway.local/core/internal/platform/legacyid"
"github.com/jackc/pgx/v5/pgxpool"
)
type record struct {
SourceSystem string `json:"source_system"`
EntityType string `json:"entity_type"`
LegacyID string `json:"legacy_id"`
Data json.RawMessage `json:"data"`
Checksum string `json:"checksum"`
}
type staged struct {
record
NewID string
}
func main() {
input := flag.String("input", "-", "JSONL file from export_legacy_data.py, or - for stdin")
dryRun := flag.Bool("dry-run", false, "validate without database writes")
flag.Parse()
reader, closeInput, err := openInput(*input)
if err != nil {
fatal(err)
}
defer closeInput()
records, sourceChecksum, counts, err := readRecords(reader)
if err != nil {
fatal(err)
}
summary := map[string]any{"records": len(records), "entities": counts, "source_checksum": sourceChecksum, "dry_run": *dryRun}
if *dryRun {
write(summary)
return
}
cfg, err := config.Load()
if err != nil {
fatal(err)
}
if cfg.Database.URL == "" {
fatal(errors.New("DATABASE_URL is required"))
}
ctx := context.Background()
pool, err := database.Open(ctx, cfg.Database)
if err != nil {
fatal(err)
}
defer pool.Close()
batchID, err := stage(ctx, pool, records, sourceChecksum, counts)
if err != nil {
fatal(err)
}
summary["batch_id"] = batchID
summary["status"] = "staged"
write(summary)
}
func openInput(path string) (io.Reader, func(), error) {
if path == "-" {
return os.Stdin, func() {}, nil
}
file, err := os.Open(path)
if err != nil {
return nil, func() {}, err
}
return file, func() { _ = file.Close() }, nil
}
func readRecords(reader io.Reader) ([]staged, string, map[string]int, error) {
scanner := bufio.NewScanner(reader)
scanner.Buffer(make([]byte, 64<<10), 32<<20)
aggregate := sha256.New()
items := []staged{}
counts := map[string]int{}
source := ""
line := 0
for scanner.Scan() {
line++
raw := bytesTrimSpace(scanner.Bytes())
if len(raw) == 0 {
continue
}
_, _ = aggregate.Write(raw)
_, _ = aggregate.Write([]byte{'\n'})
var item record
if err := json.Unmarshal(raw, &item); err != nil {
return nil, "", nil, fmt.Errorf("line %d: %w", line, err)
}
item.SourceSystem = strings.TrimSpace(item.SourceSystem)
item.EntityType = strings.TrimSpace(item.EntityType)
item.LegacyID = strings.TrimSpace(item.LegacyID)
if source == "" {
source = item.SourceSystem
}
if item.SourceSystem != source {
return nil, "", nil, fmt.Errorf("line %d: mixed source systems", line)
}
canonical, _ := json.Marshal(map[string]any{"source_system": item.SourceSystem, "entity_type": item.EntityType, "legacy_id": item.LegacyID, "data": item.Data})
digest := sha256.Sum256(canonical)
if hex.EncodeToString(digest[:]) != item.Checksum {
return nil, "", nil, fmt.Errorf("line %d: checksum mismatch", line)
}
newID, err := legacyid.UUID(item.SourceSystem, item.EntityType, item.LegacyID)
if err != nil {
return nil, "", nil, fmt.Errorf("line %d: %w", line, err)
}
items = append(items, staged{record: item, NewID: newID})
counts[item.EntityType]++
}
if err := scanner.Err(); err != nil {
return nil, "", nil, err
}
if len(items) == 0 {
return nil, "", nil, errors.New("input contains no records")
}
return items, hex.EncodeToString(aggregate.Sum(nil)), counts, nil
}
func stage(ctx context.Context, pool *pgxpool.Pool, items []staged, sourceChecksum string, counts map[string]int) (string, error) {
batchID, err := platformid.NewUUID()
if err != nil {
return "", err
}
tx, err := pool.Begin(ctx)
if err != nil {
return "", err
}
defer func() { _ = tx.Rollback(ctx) }()
countsJSON, _ := json.Marshal(counts)
source := items[0].SourceSystem
err = tx.QueryRow(ctx, `INSERT INTO gateway.legacy_import_batches(id,source_system,source_checksum,status,record_count,entity_counts) VALUES($1,$2,$3,'staged',$4,$5) ON CONFLICT(source_system,source_checksum) DO UPDATE SET source_checksum=excluded.source_checksum RETURNING id::text`, batchID, source, sourceChecksum, len(items), countsJSON).Scan(&batchID)
if err != nil {
return "", err
}
for _, item := range items {
tag, insertErr := tx.Exec(ctx, `INSERT INTO gateway.legacy_import_records(first_seen_batch_id,source_system,entity_type,legacy_id,new_id,payload,payload_checksum) VALUES($1,$2,$3,$4,$5,$6,$7) ON CONFLICT(source_system,entity_type,legacy_id) DO UPDATE SET payload_checksum=gateway.legacy_import_records.payload_checksum WHERE gateway.legacy_import_records.payload_checksum=excluded.payload_checksum`, batchID, item.SourceSystem, item.EntityType, item.LegacyID, item.NewID, item.Data, item.Checksum)
if insertErr != nil {
return "", insertErr
}
if tag.RowsAffected() == 0 {
return "", fmt.Errorf("legacy record changed since an earlier import: %s/%s", item.EntityType, item.LegacyID)
}
if _, insertErr = tx.Exec(ctx, `INSERT INTO gateway.legacy_import_batch_records(batch_id,source_system,entity_type,legacy_id,payload_checksum) VALUES($1,$2,$3,$4,$5) ON CONFLICT DO NOTHING`, batchID, item.SourceSystem, item.EntityType, item.LegacyID, item.Checksum); insertErr != nil {
return "", insertErr
}
metadata, _ := json.Marshal(map[string]any{"batch_id": batchID, "payload_checksum": item.Checksum})
if _, insertErr = tx.Exec(ctx, `INSERT INTO gateway.legacy_id_mappings(source_system,entity_type,legacy_id,new_id,metadata) VALUES($1,$2,$3,$4,$5) ON CONFLICT(source_system,entity_type,legacy_id) DO UPDATE SET metadata=excluded.metadata WHERE gateway.legacy_id_mappings.new_id=excluded.new_id`, item.SourceSystem, item.EntityType, item.LegacyID, item.NewID, metadata); insertErr != nil {
return "", insertErr
}
}
if err = tx.Commit(ctx); err != nil {
return "", err
}
return batchID, nil
}
func bytesTrimSpace(value []byte) []byte { return []byte(strings.TrimSpace(string(value))) }
func write(value any) {
encoded, _ := json.MarshalIndent(value, "", " ")
fmt.Println(string(encoded))
}
func fatal(err error) { fmt.Fprintln(os.Stderr, err); os.Exit(1) }
+27
View File
@@ -0,0 +1,27 @@
package main
import (
"bytes"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"testing"
)
func TestReadRecordsValidatesChecksumAndDeterministicID(t *testing.T) {
data := json.RawMessage(`{"id":7,"name":"研发"}`)
canonical, _ := json.Marshal(map[string]any{"source_system": "python-gateway", "entity_type": "departments", "legacy_id": "7", "data": data})
digest := sha256.Sum256(canonical)
line, _ := json.Marshal(record{SourceSystem: "python-gateway", EntityType: "departments", LegacyID: "7", Data: data, Checksum: hex.EncodeToString(digest[:])})
items, source, counts, err := readRecords(bytes.NewReader(append(line, '\n')))
if err != nil {
t.Fatal(err)
}
if len(items) != 1 || items[0].NewID == "" || source == "" || counts["departments"] != 1 {
t.Fatalf("unexpected result: %#v %s %#v", items, source, counts)
}
line[len(line)-2] ^= 1
if _, _, _, err = readRecords(bytes.NewReader(append(line, '\n'))); err == nil {
t.Fatal("expected tamper detection")
}
}
+137
View File
@@ -0,0 +1,137 @@
package main
import (
"bytes"
"context"
"encoding/json"
"flag"
"fmt"
"io"
"net/http"
"os"
"sort"
"strings"
"sync"
"sync/atomic"
"time"
)
type sample struct {
latency time.Duration
status int
err string
}
type report struct {
Requests int `json:"requests"`
Success int `json:"success"`
Errors int `json:"errors"`
ErrorRate float64 `json:"error_rate"`
RequestsPerSecond float64 `json:"requests_per_second"`
P50MS float64 `json:"p50_ms"`
P95MS float64 `json:"p95_ms"`
P99MS float64 `json:"p99_ms"`
DurationSeconds float64 `json:"duration_seconds"`
StatusCounts map[int]int `json:"status_counts"`
}
func main() {
base := flag.String("url", "http://127.0.0.1:8080", "gateway base URL")
path := flag.String("path", "/v1/chat/completions", "request path")
key := flag.String("api-key", os.Getenv("GATEWAY_LOADTEST_API_KEY"), "dedicated load-test API key")
model := flag.String("model", "test-model", "model or model alias")
concurrency := flag.Int("concurrency", 16, "concurrent workers")
duration := flag.Duration("duration", 30*time.Second, "test duration")
timeout := flag.Duration("timeout", 60*time.Second, "per-request timeout")
maxError := flag.Float64("max-error-rate", 0.01, "failure threshold")
maxP95 := flag.Duration("max-p95", 2*time.Second, "p95 latency threshold")
flag.Parse()
if *concurrency < 1 || *concurrency > 2000 || *duration < time.Second || *key == "" {
fmt.Fprintln(os.Stderr, "invalid arguments: api-key, positive duration and concurrency 1..2000 are required")
os.Exit(2)
}
target := strings.TrimRight(*base, "/") + *path
payload, _ := json.Marshal(map[string]any{"model": *model, "messages": []map[string]string{{"role": "user", "content": "Reply with OK."}}, "stream": false, "max_tokens": 8})
client := &http.Client{Timeout: *timeout, Transport: &http.Transport{MaxIdleConns: *concurrency * 2, MaxIdleConnsPerHost: *concurrency, IdleConnTimeout: 90 * time.Second, TLSHandshakeTimeout: 5 * time.Second}}
ctx, cancel := context.WithTimeout(context.Background(), *duration)
defer cancel()
started := time.Now()
results := make(chan sample, *concurrency*4)
var sequence atomic.Uint64
var workers sync.WaitGroup
for worker := 0; worker < *concurrency; worker++ {
workers.Add(1)
go func() {
defer workers.Done()
for ctx.Err() == nil {
number := sequence.Add(1)
request, err := http.NewRequestWithContext(ctx, http.MethodPost, target, bytes.NewReader(payload))
if err != nil {
results <- sample{err: err.Error()}
continue
}
request.Header.Set("Authorization", "Bearer "+*key)
request.Header.Set("Content-Type", "application/json")
request.Header.Set("X-Request-ID", fmt.Sprintf("load-%d", number))
request.Header.Set("Idempotency-Key", fmt.Sprintf("load-%d", number))
begin := time.Now()
response, err := client.Do(request)
elapsed := time.Since(begin)
if err != nil {
if ctx.Err() == nil {
results <- sample{latency: elapsed, err: err.Error()}
}
continue
}
_, readErr := io.Copy(io.Discard, io.LimitReader(response.Body, 2<<20))
response.Body.Close()
entry := sample{latency: elapsed, status: response.StatusCode}
if readErr != nil {
entry.err = readErr.Error()
}
results <- entry
}
}()
}
go func() { workers.Wait(); close(results) }()
samples := []sample{}
for result := range results {
samples = append(samples, result)
}
elapsed := time.Since(started)
summary := summarize(samples, elapsed)
encoded, _ := json.MarshalIndent(summary, "", " ")
fmt.Println(string(encoded))
if summary.ErrorRate > *maxError || time.Duration(summary.P95MS*float64(time.Millisecond)) > *maxP95 {
os.Exit(1)
}
}
func summarize(samples []sample, elapsed time.Duration) report {
summary := report{Requests: len(samples), StatusCounts: map[int]int{}, DurationSeconds: elapsed.Seconds()}
latencies := make([]time.Duration, 0, len(samples))
for _, item := range samples {
summary.StatusCounts[item.status]++
latencies = append(latencies, item.latency)
if item.err == "" && item.status >= 200 && item.status < 300 {
summary.Success++
} else {
summary.Errors++
}
}
sort.Slice(latencies, func(i, j int) bool { return latencies[i] < latencies[j] })
if summary.Requests > 0 {
summary.ErrorRate = float64(summary.Errors) / float64(summary.Requests)
summary.RequestsPerSecond = float64(summary.Requests) / elapsed.Seconds()
summary.P50MS = percentile(latencies, .50)
summary.P95MS = percentile(latencies, .95)
summary.P99MS = percentile(latencies, .99)
}
return summary
}
func percentile(values []time.Duration, p float64) float64 {
if len(values) == 0 {
return 0
}
index := int(float64(len(values)-1) * p)
return float64(values[index]) / float64(time.Millisecond)
}
+16
View File
@@ -0,0 +1,16 @@
package main
import (
"testing"
"time"
)
func TestSummarize(t *testing.T) {
summary := summarize([]sample{{latency: 10 * time.Millisecond, status: 200}, {latency: 20 * time.Millisecond, status: 200}, {latency: 100 * time.Millisecond, status: 500}}, time.Second)
if summary.Requests != 3 || summary.Success != 2 || summary.Errors != 1 {
t.Fatalf("unexpected counts: %#v", summary)
}
if summary.P95MS != 20 {
t.Fatalf("unexpected p95: %v", summary.P95MS)
}
}
+52
View File
@@ -0,0 +1,52 @@
package main
import (
"context"
"log/slog"
"os"
"os/signal"
"syscall"
"time"
"aigateway.local/core/internal/audit"
"aigateway.local/core/internal/platform/config"
"aigateway.local/core/internal/platform/database"
)
func main() {
logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo}))
cfg, err := config.Load()
if err != nil {
logger.Error("invalid configuration", "error", err)
os.Exit(1)
}
if cfg.Database.URL == "" {
logger.Error("DATABASE_URL is required")
os.Exit(1)
}
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
defer stop()
db, err := database.Open(ctx, cfg.Database)
if err != nil {
logger.Error("database initialization failed", "error", err)
os.Exit(1)
}
defer db.Close()
maintenance := audit.NewMaintenance(db, cfg.Audit.Retention, cfg.Audit.UsageRetention, cfg.Audit.PartitionMonthsAhead)
ticker := time.NewTicker(cfg.Audit.MaintenanceInterval)
defer ticker.Stop()
for {
result, runErr := maintenance.Run(ctx, time.Now())
if runErr != nil && ctx.Err() == nil {
logger.Error("audit maintenance failed", "error", runErr)
} else if runErr == nil {
logger.Info("audit maintenance complete", "created_partitions", result.CreatedPartitions, "dropped_partitions", result.DroppedPartitions, "deleted_audit_rows", result.DeletedAuditRows, "deleted_usage_rows", result.DeletedUsageRows)
}
select {
case <-ctx.Done():
logger.Info("maintenance worker stopped")
return
case <-ticker.C:
}
}
}
+44
View File
@@ -0,0 +1,44 @@
package main
import (
"context"
"flag"
"log/slog"
"os"
"aigateway.local/core/internal/platform/config"
"aigateway.local/core/internal/platform/database"
"aigateway.local/core/internal/platform/migrate"
)
func main() {
directory := flag.String("dir", "migrations", "directory containing ordered SQL migrations")
flag.Parse()
logger := slog.New(slog.NewJSONHandler(os.Stdout, nil))
cfg, err := config.Load()
if err != nil {
logger.Error("invalid configuration", "error", err)
os.Exit(1)
}
if cfg.Database.URL == "" {
logger.Error("DATABASE_URL is required by the migrator")
os.Exit(1)
}
ctx := context.Background()
pool, err := database.Open(ctx, cfg.Database)
if err != nil {
logger.Error("database initialization failed", "error", err)
os.Exit(1)
}
defer pool.Close()
migrations, err := migrate.Load(*directory)
if err != nil {
logger.Error("migration loading failed", "error", err)
os.Exit(1)
}
if err := migrate.Apply(ctx, pool, migrations); err != nil {
logger.Error("migration failed", "error", err)
os.Exit(1)
}
logger.Info("migrations complete", "count", len(migrations))
}
+61
View File
@@ -0,0 +1,61 @@
package main
import (
"context"
"log/slog"
"os"
"os/signal"
"syscall"
"aigateway.local/core/internal/platform/cache"
"aigateway.local/core/internal/platform/config"
"aigateway.local/core/internal/platform/cryptox"
"aigateway.local/core/internal/platform/database"
platformid "aigateway.local/core/internal/platform/id"
"aigateway.local/core/internal/workbench"
)
func main() {
logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo}))
cfg, err := config.Load()
if err != nil {
logger.Error("invalid configuration", "error", err)
os.Exit(1)
}
if cfg.Database.URL == "" || cfg.Redis.CriticalURL == "" {
logger.Error("DATABASE_URL and REDIS_CRITICAL_URL are required")
os.Exit(1)
}
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
defer stop()
db, err := database.Open(ctx, cfg.Database)
if err != nil {
logger.Error("database initialization failed", "error", err)
os.Exit(1)
}
defer db.Close()
client, err := cache.Open(cfg.Redis.CriticalURL)
if err != nil {
logger.Error("redis initialization failed", "error", err)
os.Exit(1)
}
defer client.Close()
cipher, err := cryptox.NewKeyring(cfg.Credentials.MasterKey, cfg.Credentials.KEKVersion, cfg.Credentials.KEKKeyring, "notification-signing-secret")
if err != nil {
logger.Error("notification cipher initialization failed", "error", err)
os.Exit(1)
}
consumer, err := platformid.NewUUID()
if err != nil {
logger.Error("consumer ID generation failed", "error", err)
os.Exit(1)
}
service := workbench.NewNotificationService(workbench.NewService(db), cipher, cfg.Credentials.AllowPrivateWebhookURL)
dispatcher := workbench.NewNotificationDispatcher(service, client, cfg.Outbox.Stream, "notification-"+consumer, logger)
logger.Info("notification worker started", "consumer", consumer, "stream", cfg.Outbox.Stream)
if err = dispatcher.Run(ctx); err != nil {
logger.Error("notification worker stopped unexpectedly", "error", err)
os.Exit(1)
}
logger.Info("notification worker stopped", "consumer", consumer)
}
+68
View File
@@ -0,0 +1,68 @@
package main
import (
"context"
"log/slog"
"os"
"os/signal"
"syscall"
"aigateway.local/core/internal/outbox"
"aigateway.local/core/internal/platform/cache"
"aigateway.local/core/internal/platform/config"
"aigateway.local/core/internal/platform/database"
platformid "aigateway.local/core/internal/platform/id"
)
func main() {
logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo}))
cfg, err := config.Load()
if err != nil {
logger.Error("invalid configuration", "error", err)
os.Exit(1)
}
if cfg.Database.URL == "" || cfg.Redis.CriticalURL == "" {
logger.Error("DATABASE_URL and REDIS_CRITICAL_URL are required")
os.Exit(1)
}
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
defer stop()
db, err := database.Open(ctx, cfg.Database)
if err != nil {
logger.Error("database initialization failed", "error", err)
os.Exit(1)
}
defer db.Close()
redisClient, err := cache.Open(cfg.Redis.CriticalURL)
if err != nil {
logger.Error("redis initialization failed", "error", err)
os.Exit(1)
}
defer redisClient.Close()
if err := db.Ping(ctx); err != nil {
logger.Error("database is unavailable", "error", err)
os.Exit(1)
}
if err := redisClient.Ping(ctx).Err(); err != nil {
logger.Error("redis is unavailable", "error", err)
os.Exit(1)
}
workerID, err := platformid.NewUUID()
if err != nil {
logger.Error("worker ID generation failed", "error", err)
os.Exit(1)
}
workerID = "outbox-" + workerID
worker := outbox.NewWorker(
outbox.NewStore(db),
outbox.NewRedisPublisher(redisClient, cfg.Outbox.Stream, cfg.Outbox.StreamMaxLength, cfg.Outbox.MarkerTTL),
outbox.WorkerConfig{WorkerID: workerID, BatchSize: cfg.Outbox.BatchSize, PollInterval: cfg.Outbox.PollInterval, Lease: cfg.Outbox.Lease, MaxAttempts: cfg.Outbox.MaxAttempts, MaxBackoff: cfg.Outbox.MaxBackoff},
logger,
)
logger.Info("outbox worker started", "worker_id", workerID, "stream", cfg.Outbox.Stream)
if err := worker.Run(ctx); err != nil {
logger.Error("outbox worker stopped unexpectedly", "error", err)
os.Exit(1)
}
logger.Info("outbox worker stopped", "worker_id", workerID)
}