package access import ( "context" "fmt" "strings" "time" "sub2api-cn-relay-manager/internal/host/sub2api" ) const ( ModeSubscription = "subscription" ModeSelfService = "self_service" gatewayCompletionRetryAttempts = 3 gatewayCompletionRetryDelay = 300 * time.Millisecond ) type SubscriptionTarget struct { UserID string DurationDays int } type ClosureRequest struct { Mode string ProbeAPIKey string Subscriptions []SubscriptionTarget GroupID string ExpectedModel string Prompt string MaxTokens int } type Host interface { EnsureSubscriptionAccess(ctx context.Context, req sub2api.EnsureSubscriptionAccessRequest) (sub2api.SubscriptionAccessRef, error) AssignSubscription(ctx context.Context, req sub2api.AssignSubscriptionRequest) (sub2api.SubscriptionRef, error) CheckGatewayAccess(ctx context.Context, req sub2api.GatewayAccessCheckRequest) (sub2api.GatewayAccessResult, error) CheckGatewayCompletion(ctx context.Context, req sub2api.GatewayCompletionCheckRequest) (sub2api.GatewayCompletionResult, error) } type Service struct { host Host } func NewService(host Host) *Service { return &Service{host: host} } func Validate(req ClosureRequest) error { switch strings.TrimSpace(req.Mode) { case ModeSubscription: if len(req.Subscriptions) == 0 { return fmt.Errorf("subscription access requires at least one subscription target") } case ModeSelfService: if strings.TrimSpace(req.ProbeAPIKey) == "" { return fmt.Errorf("self_service access requires probe api key") } default: return fmt.Errorf("unsupported access mode %q", req.Mode) } if strings.TrimSpace(req.Mode) != ModeSubscription && strings.TrimSpace(req.ProbeAPIKey) == "" { return fmt.Errorf("access probe api key is required to verify gateway closure") } return nil } func (s *Service) Close(ctx context.Context, req ClosureRequest) (sub2api.GatewayAccessResult, error) { if s == nil || s.host == nil { return sub2api.GatewayAccessResult{}, fmt.Errorf("access host is required") } if err := Validate(req); err != nil { return sub2api.GatewayAccessResult{}, err } probeAPIKey := strings.TrimSpace(req.ProbeAPIKey) if strings.TrimSpace(req.Mode) == ModeSubscription { for _, target := range req.Subscriptions { resolvedTarget := target.UserID accessRef, err := s.host.EnsureSubscriptionAccess(ctx, sub2api.EnsureSubscriptionAccessRequest{UserSelector: target.UserID, GroupID: req.GroupID}) if err != nil { return sub2api.GatewayAccessResult{}, fmt.Errorf("ensure subscription access for %s: %w", target.UserID, err) } if strings.TrimSpace(accessRef.UserID) != "" { resolvedTarget = accessRef.UserID } if strings.TrimSpace(accessRef.APIKey) != "" { probeAPIKey = strings.TrimSpace(accessRef.APIKey) } if _, err := s.host.AssignSubscription(ctx, sub2api.AssignSubscriptionRequest{UserID: resolvedTarget, GroupID: req.GroupID, DurationDays: target.DurationDays}); err != nil { return sub2api.GatewayAccessResult{}, fmt.Errorf("assign subscription for %s: %w", target.UserID, err) } } } if probeAPIKey == "" { return sub2api.GatewayAccessResult{}, fmt.Errorf("access probe api key is required to verify gateway closure") } result, err := s.host.CheckGatewayAccess(ctx, sub2api.GatewayAccessCheckRequest{APIKey: probeAPIKey, ExpectedModel: req.ExpectedModel}) if err != nil { return sub2api.GatewayAccessResult{}, fmt.Errorf("check gateway access: %w", err) } if result.OK && result.HasExpectedModel && strings.TrimSpace(req.ExpectedModel) != "" { completion, err := s.checkGatewayCompletionWithRetry(ctx, sub2api.GatewayCompletionCheckRequest{ APIKey: probeAPIKey, Model: req.ExpectedModel, Prompt: req.Prompt, MaxTokens: req.MaxTokens, }) if err != nil { return sub2api.GatewayAccessResult{}, fmt.Errorf("check gateway completion: %w", err) } result.CompletionOK = completion.OK result.CompletionStatus = completion.StatusCode result.CompletionType = completion.ContentType result.CompletionBody = completion.BodyPreview } return result, nil } func (s *Service) checkGatewayCompletionWithRetry(ctx context.Context, req sub2api.GatewayCompletionCheckRequest) (sub2api.GatewayCompletionResult, error) { var last sub2api.GatewayCompletionResult for attempt := 1; attempt <= gatewayCompletionRetryAttempts; attempt++ { completion, err := s.host.CheckGatewayCompletion(ctx, req) if err != nil { return sub2api.GatewayCompletionResult{}, err } last = completion if completion.OK || !isTransientGatewayCompletionFailure(completion) || attempt == gatewayCompletionRetryAttempts { return completion, nil } timer := time.NewTimer(gatewayCompletionRetryDelay) select { case <-ctx.Done(): timer.Stop() return last, ctx.Err() case <-timer.C: } } return last, nil } func isTransientGatewayCompletionFailure(result sub2api.GatewayCompletionResult) bool { if result.OK { return false } if result.StatusCode != 0 && result.StatusCode != 429 && result.StatusCode != 502 && result.StatusCode != 503 && result.StatusCode != 504 { return false } body := strings.ToLower(strings.TrimSpace(result.BodyPreview)) return strings.Contains(body, "service temporarily unavailable") || strings.Contains(body, "no available accounts") || strings.Contains(body, "temporar") || strings.Contains(body, "try again") }