sub2api-cn-relay-manager/internal/access/closure.go

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")
}