package main import ( "context" "crypto/rand" "encoding/hex" "encoding/json" "errors" "flag" "fmt" "os" "path/filepath" "regexp" "strings" "time" "github.com/easyai/easyai-ai-gateway/apps/api/internal/acceptancesnapshot" "github.com/easyai/easyai-ai-gateway/apps/api/internal/auth" "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" ) var ( fullSHAPattern = regexp.MustCompile(`^[0-9a-f]{40}$`) digestPattern = regexp.MustCompile(`^sha256:[0-9a-f]{64}$`) ) type options struct { clusterID string releaseSHA string apiImageDigest string workerDigest string emulatorBaseURL string callbackURL string output string identityShards int } type runtimeFile struct { SchemaVersion string `json:"schemaVersion"` LocalClusterID string `json:"localClusterId"` RunID string `json:"runId"` RunToken string `json:"runToken"` APIKeys []string `json:"apiKeys"` Participants []runtimeParticipant `json:"participants"` GeminiModel string `json:"geminiModel"` VideoModel string `json:"videoModel"` EmulatorBaseURL string `json:"emulatorBaseUrl"` CallbackURL string `json:"callbackUrl"` ReleaseSHA string `json:"releaseSha"` APIImageDigest string `json:"apiImageDigest"` WorkerImageDigest string `json:"workerImageDigest"` SnapshotConfigHash string `json:"snapshotConfigHash"` SnapshotSHA256 string `json:"snapshotSha256"` CreatedAt time.Time `json:"createdAt"` } type runtimeParticipant struct { APIKeyID string `json:"apiKeyId"` UserID string `json:"userId"` } func main() { opts, err := parseOptions() if err != nil { fmt.Fprintln(os.Stderr, "acceptance bootstrap:", err) os.Exit(64) } ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) defer cancel() if err := run(ctx, opts); err != nil { fmt.Fprintln(os.Stderr, "acceptance bootstrap:", err) os.Exit(1) } } func parseOptions() (options, error) { var opts options flag.StringVar(&opts.clusterID, "local-cluster-id", "", "required local cluster marker") flag.StringVar(&opts.releaseSHA, "release-sha", "", "full source Git SHA") flag.StringVar(&opts.apiImageDigest, "api-image-digest", "", "immutable API image digest") flag.StringVar(&opts.workerDigest, "worker-image-digest", "", "immutable Worker image digest") flag.StringVar(&opts.emulatorBaseURL, "emulator-base-url", "", "in-cluster protocol emulator URL") flag.StringVar(&opts.callbackURL, "callback-url", "", "in-cluster callback collector URL") flag.StringVar(&opts.output, "output", "", "private runtime output file") flag.IntVar(&opts.identityShards, "identity-shards", 32, "isolated acceptance identities") flag.Parse() opts.clusterID = strings.TrimSpace(opts.clusterID) opts.releaseSHA = strings.ToLower(strings.TrimSpace(opts.releaseSHA)) opts.apiImageDigest = strings.ToLower(strings.TrimSpace(opts.apiImageDigest)) opts.workerDigest = strings.ToLower(strings.TrimSpace(opts.workerDigest)) opts.emulatorBaseURL = strings.TrimRight(strings.TrimSpace(opts.emulatorBaseURL), "/") opts.callbackURL = strings.TrimSpace(opts.callbackURL) if opts.clusterID == "" || !fullSHAPattern.MatchString(opts.releaseSHA) || !digestPattern.MatchString(opts.apiImageDigest) || !digestPattern.MatchString(opts.workerDigest) { return options{}, errors.New("local cluster ID, full release SHA, and immutable image digests are required") } if opts.emulatorBaseURL == "" || opts.callbackURL == "" { return options{}, errors.New("emulator and callback URLs are required") } if opts.identityShards < 1 || opts.identityShards > 128 { return options{}, errors.New("identity shards must be between 1 and 128") } if strings.TrimSpace(opts.output) == "" { return options{}, errors.New("private runtime output path is required") } if flag.NArg() != 0 { return options{}, errors.New("unexpected positional arguments") } return opts, nil } func run(ctx context.Context, opts options) error { databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_DATABASE_URL")) if databaseURL == "" { return errors.New("AI_GATEWAY_DATABASE_URL is required") } database, err := store.ConnectWithMaxConns(ctx, databaseURL, 4) if err != nil { return err } defer database.Close() if err := verifyLocalDatabase(ctx, database, opts.clusterID); err != nil { return err } snapshotConfigHash, snapshotSHA, err := importedSnapshot(ctx, database) if err != nil { return err } groupID, tenantID, err := ensureAcceptanceIdentityDomain(ctx, database) if err != nil { return err } participants := make([]runtimeParticipant, 0, opts.identityShards) apiKeys := make([]string, 0, opts.identityShards) for ordinal := 0; ordinal < opts.identityShards; ordinal++ { userID, err := ensureAcceptanceUser(ctx, database, groupID, tenantID, ordinal) if err != nil { return err } keyID, secret, err := ensureAcceptanceAPIKey(ctx, database, userID, groupID, tenantID, ordinal) if err != nil { return err } if _, err := database.SetUserWalletBalance(ctx, store.WalletBalanceAdjustmentInput{ GatewayUserID: userID, Currency: "resource", BalanceText: "1000000000", Reason: "local acceptance isolated wallet", }); err != nil && !errors.Is(err, store.ErrWalletBalanceUnchanged) { return err } participants = append(participants, runtimeParticipant{APIKeyID: keyID, UserID: userID}) apiKeys = append(apiKeys, secret) } if err := ensureAcceptanceAccessRules(ctx, database, groupID); err != nil { return err } geminiModel, videoModel, err := selectedModels(ctx, database) if err != nil { return err } runToken, err := randomToken() if err != nil { return err } run, err := database.CreateAcceptanceRun(ctx, store.CreateAcceptanceRunInput{ ReleaseSHA: opts.releaseSHA, APIImageDigest: opts.apiImageDigest, WorkerImageDigest: opts.workerDigest, APIKeyID: participants[0].APIKeyID, UserID: participants[0].UserID, Token: runToken, EmulatorBaseURL: opts.emulatorBaseURL, CallbackURL: opts.callbackURL, CapacityProfile: "P24", Config: map[string]any{ "workloads": []any{"gemini_image_edit", "multi_reference_video"}, "participants": participants, "localClusterId": opts.clusterID, "snapshotConfigHash": snapshotConfigHash, "snapshotSha256": snapshotSHA, }, }) if err != nil { return err } if _, err := database.ActivateAcceptanceRun(ctx, run.ID); err != nil { return err } outputPath, err := privateOutputPath(opts.output) if err != nil { return err } payload, err := json.MarshalIndent(runtimeFile{ SchemaVersion: "acceptance-runtime/v1", LocalClusterID: opts.clusterID, RunID: run.ID, RunToken: runToken, APIKeys: apiKeys, Participants: participants, GeminiModel: geminiModel, VideoModel: videoModel, EmulatorBaseURL: opts.emulatorBaseURL, CallbackURL: opts.callbackURL, ReleaseSHA: opts.releaseSHA, APIImageDigest: opts.apiImageDigest, WorkerImageDigest: opts.workerDigest, SnapshotConfigHash: snapshotConfigHash, SnapshotSHA256: snapshotSHA, CreatedAt: time.Now().UTC(), }, "", " ") if err != nil { return err } if err := os.WriteFile(outputPath, append(payload, '\n'), 0o600); err != nil { return err } fmt.Printf( "acceptance_bootstrap=PASS cluster_id=%s run_id=%s identities=%d config_hash=%s runtime_file=%s\n", opts.clusterID, run.ID, len(participants), snapshotConfigHash, outputPath, ) return nil } func verifyLocalDatabase(ctx context.Context, database *store.Store, clusterID string) error { var marker string err := database.Pool().QueryRow(ctx, ` SELECT COALESCE(value->>'clusterId', '') FROM system_settings WHERE setting_key = $1`, acceptancesnapshot.LocalClusterSettingKey).Scan(&marker) if err != nil { return fmt.Errorf("read local acceptance cluster marker: %w", err) } if marker != clusterID { return errors.New("refusing bootstrap: local cluster marker mismatch") } mode, err := database.GetGatewayTrafficMode(ctx) if err != nil { return err } if mode.Mode != "live" { return fmt.Errorf("refusing bootstrap while traffic mode is %s", mode.Mode) } return nil } func importedSnapshot(ctx context.Context, database *store.Store) (string, string, error) { var configHash, snapshotSHA string err := database.Pool().QueryRow(ctx, ` SELECT COALESCE(value->>'configHash', ''), COALESCE(value->>'snapshotSha256', '') FROM system_settings WHERE setting_key = 'acceptance_snapshot'`).Scan(&configHash, &snapshotSHA) if err != nil { return "", "", fmt.Errorf("read imported acceptance snapshot: %w", err) } if len(configHash) != 64 || len(snapshotSHA) != 64 { return "", "", errors.New("imported acceptance snapshot hashes are invalid") } return configHash, snapshotSHA, nil } func ensureAcceptanceIdentityDomain(ctx context.Context, database *store.Store) (string, string, error) { var groupID string err := database.Pool().QueryRow(ctx, ` INSERT INTO gateway_user_groups ( group_key, name, description, source, priority, recharge_discount_policy, billing_discount_policy, rate_limit_policy, quota_policy, metadata, status ) VALUES ( 'local-acceptance', 'Local Acceptance', 'Isolated local homologous acceptance identities', 'gateway', 1, '{"discountFactor":1}'::jsonb, '{"discountFactor":1}'::jsonb, '{"rules":[]}'::jsonb, '{}'::jsonb, '{"purpose":"local_acceptance","isolated":true}'::jsonb, 'active' ) ON CONFLICT (group_key) DO UPDATE SET name = EXCLUDED.name, description = EXCLUDED.description, priority = EXCLUDED.priority, rate_limit_policy = EXCLUDED.rate_limit_policy, quota_policy = EXCLUDED.quota_policy, metadata = EXCLUDED.metadata, status = 'active', updated_at = now() RETURNING id::text`).Scan(&groupID) if err != nil { return "", "", err } var tenantID string err = database.Pool().QueryRow(ctx, ` INSERT INTO gateway_tenants ( tenant_key, source, external_tenant_id, name, default_user_group_id, billing_profile, rate_limit_policy, auth_policy, metadata, status ) VALUES ( 'local-acceptance', 'gateway', 'local-acceptance', 'Local Acceptance', $1::uuid, '{}'::jsonb, '{"rules":[]}'::jsonb, '{}'::jsonb, '{"purpose":"local_acceptance","isolated":true}'::jsonb, 'active' ) ON CONFLICT (tenant_key) DO UPDATE SET default_user_group_id = EXCLUDED.default_user_group_id, metadata = EXCLUDED.metadata, status = 'active', updated_at = now() RETURNING id::text`, groupID).Scan(&tenantID) return groupID, tenantID, err } func ensureAcceptanceUser( ctx context.Context, database *store.Store, groupID string, tenantID string, ordinal int, ) (string, error) { userKey := fmt.Sprintf("local-acceptance-%03d", ordinal) roles := `["user"]` if ordinal == 0 { roles = `["manager"]` } var userID string err := database.Pool().QueryRow(ctx, ` INSERT INTO gateway_users ( user_key, source, external_user_id, username, display_name, gateway_tenant_id, tenant_id, tenant_key, default_user_group_id, roles, auth_profile, metadata, status ) VALUES ( $1, 'gateway', $1, $1, $2, $3::uuid, 'local-acceptance', 'local-acceptance', $4::uuid, $5::jsonb, '{}'::jsonb, jsonb_build_object('purpose','local_acceptance','ordinal',$6,'isolated',true), 'active' ) ON CONFLICT (user_key) DO UPDATE SET gateway_tenant_id = EXCLUDED.gateway_tenant_id, tenant_id = EXCLUDED.tenant_id, tenant_key = EXCLUDED.tenant_key, default_user_group_id = EXCLUDED.default_user_group_id, roles = EXCLUDED.roles, metadata = EXCLUDED.metadata, status = 'active', deleted_at = NULL, updated_at = now() RETURNING id::text`, userKey, fmt.Sprintf("Local Acceptance %03d", ordinal), tenantID, groupID, roles, ordinal, ).Scan(&userID) return userID, err } func ensureAcceptanceAPIKey( ctx context.Context, database *store.Store, userID string, groupID string, tenantID string, ordinal int, ) (string, string, error) { name := fmt.Sprintf("Local Acceptance %03d", ordinal) var keyID, secret string err := database.Pool().QueryRow(ctx, ` SELECT id::text, COALESCE(key_secret, '') FROM gateway_api_keys WHERE gateway_user_id = $1::uuid AND name = $2 AND status = 'active' AND deleted_at IS NULL AND COALESCE(key_secret, '') <> '' ORDER BY created_at LIMIT 1`, userID, name).Scan(&keyID, &secret) if err == nil { return keyID, secret, nil } if !store.IsNotFound(err) { return "", "", err } created, err := database.CreateAPIKey(ctx, store.CreateAPIKeyInput{ Name: name, Scopes: []string{"image", "video"}, }, &auth.User{ ID: userID, Username: fmt.Sprintf("local-acceptance-%03d", ordinal), GatewayUserID: userID, GatewayTenantID: tenantID, TenantID: "local-acceptance", TenantKey: "local-acceptance", UserGroupID: groupID, Source: "gateway", }) if err != nil { return "", "", err } if _, err := database.Pool().Exec(ctx, ` UPDATE gateway_api_keys SET user_group_id = $2::uuid, rate_limit_policy = '{"rules":[]}'::jsonb, quota_policy = '{}'::jsonb, updated_at = now() WHERE id = $1::uuid`, created.APIKey.ID, groupID); err != nil { return "", "", err } return created.APIKey.ID, created.Secret, nil } func ensureAcceptanceAccessRules(ctx context.Context, database *store.Store, groupID string) error { _, err := database.Pool().Exec(ctx, ` WITH selected AS ( SELECT 'platform'::text AS resource_type, platform.id AS resource_id FROM integration_platforms platform WHERE COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false) UNION SELECT 'platform_model', model.id FROM platform_models model JOIN integration_platforms platform ON platform.id = model.platform_id WHERE COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false) UNION SELECT 'base_model', model.base_model_id FROM platform_models model JOIN integration_platforms platform ON platform.id = model.platform_id WHERE COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false) AND model.base_model_id IS NOT NULL ) INSERT INTO gateway_access_rules ( subject_type, subject_id, resource_type, resource_id, effect, priority, min_permission_level, conditions, metadata, status ) SELECT 'user_group', $1::uuid, selected.resource_type, selected.resource_id, 'allow', 1, 0, '{}'::jsonb, '{"purpose":"local_acceptance"}'::jsonb, 'active' FROM selected ON CONFLICT (subject_type, subject_id, resource_type, resource_id, effect) DO UPDATE SET status = 'active', priority = 1, updated_at = now()`, groupID) return err } func selectedModels(ctx context.Context, database *store.Store) (string, string, error) { var gemini, video string err := database.Pool().QueryRow(ctx, ` SELECT base_model.invocation_name FROM platform_models model JOIN integration_platforms platform ON platform.id = model.platform_id JOIN base_model_catalog base_model ON base_model.id = model.base_model_id WHERE platform.status = 'enabled' AND COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false) AND model.enabled = true AND model.model_type @> '["image_edit"]'::jsonb ORDER BY platform.priority, model.created_at LIMIT 1`).Scan(&gemini) if err != nil { return "", "", err } err = database.Pool().QueryRow(ctx, ` SELECT base_model.invocation_name FROM platform_models model JOIN integration_platforms platform ON platform.id = model.platform_id JOIN base_model_catalog base_model ON base_model.id = model.base_model_id WHERE platform.status = 'enabled' AND COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false) AND model.enabled = true AND model.model_type @> '["omni_video"]'::jsonb ORDER BY platform.priority, model.created_at LIMIT 1`).Scan(&video) return gemini, video, err } func randomToken() (string, error) { payload := make([]byte, 32) if _, err := rand.Read(payload); err != nil { return "", err } return hex.EncodeToString(payload), nil } func privateOutputPath(path string) (string, error) { absolute, err := filepath.Abs(strings.TrimSpace(path)) if err != nil { return "", err } if info, err := os.Lstat(absolute); err == nil { if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 { return "", errors.New("runtime output must be a regular non-symlink file") } } else if !os.IsNotExist(err) { return "", err } if err := os.MkdirAll(filepath.Dir(absolute), 0o700); err != nil { return "", err } return absolute, nil }