sub2api-cn-relay-manager/internal/routing/sticky.go

package routing

import (
	"context"
	"fmt"
	"strings"
	"time"
)

const (
	RuntimeBackendMemory = "memory"
	RuntimeBackendRedis  = "redis"

	StickyScopeConversation = "conversation"
	StickyScopeSession      = "session"
	StickyScopeUser         = "user"
)

type StickyBinding struct {
	LogicalGroupID string `json:"logical_group_id"`
	PublicModel    string `json:"public_model"`
	RouteID        string `json:"route_id"`
	ShadowGroupID  string `json:"shadow_group_id"`
	BoundAt        string `json:"bound_at,omitempty"`
	ExpiresAt      string `json:"expires_at,omitempty"`
}

type RouteFailureState struct {
	RouteID        string `json:"route_id"`
	FailureCount   int    `json:"failure_count"`
	LastErrorClass string `json:"last_error_class,omitempty"`
	LastFailureAt  string `json:"last_failure_at,omitempty"`
	ExpiresAt      string `json:"expires_at,omitempty"`
}

type RouteCooldownState struct {
	RouteID   string `json:"route_id"`
	Reason    string `json:"reason,omitempty"`
	Until     string `json:"until,omitempty"`
	ExpiresAt string `json:"expires_at,omitempty"`
}

type StickyStore interface {
	Get(ctx context.Context, key string) (StickyBinding, bool, error)
	Set(ctx context.Context, key string, binding StickyBinding, ttl time.Duration) error
	Delete(ctx context.Context, key string) error

	GetRouteFailure(ctx context.Context, routeID string) (RouteFailureState, bool, error)
	SetRouteFailure(ctx context.Context, routeID string, state RouteFailureState, ttl time.Duration) error
	ClearRouteFailure(ctx context.Context, routeID string) error

	GetCooldown(ctx context.Context, routeID string) (RouteCooldownState, bool, error)
	SetCooldown(ctx context.Context, routeID string, state RouteCooldownState, ttl time.Duration) error
	ClearCooldown(ctx context.Context, routeID string) error
}

type StoreConfig struct {
	Backend string
	Redis   RedisConfig
}

func NormalizeRuntimeBackend(backend string) (string, error) {
	switch strings.ToLower(strings.TrimSpace(backend)) {
	case "", RuntimeBackendMemory:
		return RuntimeBackendMemory, nil
	case RuntimeBackendRedis:
		return RuntimeBackendRedis, nil
	default:
		return "", fmt.Errorf("unsupported route runtime backend %q", backend)
	}
}

func BuildStickyKey(scope, logicalGroupID, publicModel, subjectID string) (string, error) {
	scope = strings.ToLower(strings.TrimSpace(scope))
	logicalGroupID = strings.TrimSpace(logicalGroupID)
	publicModel = strings.TrimSpace(publicModel)
	subjectID = strings.TrimSpace(subjectID)

	switch {
	case logicalGroupID == "":
		return "", fmt.Errorf("logical_group_id is required")
	case publicModel == "":
		return "", fmt.Errorf("public_model is required")
	case subjectID == "":
		return "", fmt.Errorf("subject_id is required")
	}

	var subjectPrefix string
	switch scope {
	case StickyScopeConversation:
		subjectPrefix = "conv"
	case StickyScopeSession:
		subjectPrefix = "sess"
	case StickyScopeUser:
		subjectPrefix = "user"
	default:
		return "", fmt.Errorf("unsupported sticky scope %q", scope)
	}

	return fmt.Sprintf("lg:%s:m:%s:%s:%s", logicalGroupID, publicModel, subjectPrefix, subjectID), nil
}

func BuildRouteFailureKey(routeID string) (string, error) {
	routeID = strings.TrimSpace(routeID)
	if routeID == "" {
		return "", fmt.Errorf("route_id is required")
	}
	return "routefail:" + routeID, nil
}

func BuildRouteCooldownKey(routeID string) (string, error) {
	routeID = strings.TrimSpace(routeID)
	if routeID == "" {
		return "", fmt.Errorf("route_id is required")
	}
	return "routecool:" + routeID, nil
}

func normalizeStickyBinding(binding StickyBinding, ttl time.Duration, now time.Time) (StickyBinding, error) {
	binding.LogicalGroupID = strings.TrimSpace(binding.LogicalGroupID)
	binding.PublicModel = strings.TrimSpace(binding.PublicModel)
	binding.RouteID = strings.TrimSpace(binding.RouteID)
	binding.ShadowGroupID = strings.TrimSpace(binding.ShadowGroupID)

	switch {
	case binding.LogicalGroupID == "":
		return StickyBinding{}, fmt.Errorf("logical_group_id is required")
	case binding.PublicModel == "":
		return StickyBinding{}, fmt.Errorf("public_model is required")
	case binding.RouteID == "":
		return StickyBinding{}, fmt.Errorf("route_id is required")
	case binding.ShadowGroupID == "":
		return StickyBinding{}, fmt.Errorf("shadow_group_id is required")
	case ttl <= 0:
		return StickyBinding{}, fmt.Errorf("ttl must be positive")
	}

	if binding.BoundAt == "" {
		binding.BoundAt = now.UTC().Format(time.RFC3339)
	}
	if binding.ExpiresAt == "" {
		binding.ExpiresAt = now.UTC().Add(ttl).Format(time.RFC3339)
	}
	return binding, nil
}

func normalizeRouteFailureState(routeID string, state RouteFailureState, ttl time.Duration, now time.Time) (RouteFailureState, error) {
	routeID = strings.TrimSpace(routeID)
	state.RouteID = strings.TrimSpace(state.RouteID)
	state.LastErrorClass = strings.TrimSpace(state.LastErrorClass)

	if state.RouteID == "" {
		state.RouteID = routeID
	}
	switch {
	case routeID == "":
		return RouteFailureState{}, fmt.Errorf("route_id is required")
	case state.RouteID != routeID:
		return RouteFailureState{}, fmt.Errorf("route_id mismatch")
	case state.FailureCount < 0:
		return RouteFailureState{}, fmt.Errorf("failure_count must be >= 0")
	case ttl <= 0:
		return RouteFailureState{}, fmt.Errorf("ttl must be positive")
	}

	if state.LastFailureAt == "" {
		state.LastFailureAt = now.UTC().Format(time.RFC3339)
	}
	if state.ExpiresAt == "" {
		state.ExpiresAt = now.UTC().Add(ttl).Format(time.RFC3339)
	}
	return state, nil
}

func normalizeRouteCooldownState(routeID string, state RouteCooldownState, ttl time.Duration, now time.Time) (RouteCooldownState, error) {
	routeID = strings.TrimSpace(routeID)
	state.RouteID = strings.TrimSpace(state.RouteID)
	state.Reason = strings.TrimSpace(state.Reason)

	if state.RouteID == "" {
		state.RouteID = routeID
	}
	switch {
	case routeID == "":
		return RouteCooldownState{}, fmt.Errorf("route_id is required")
	case state.RouteID != routeID:
		return RouteCooldownState{}, fmt.Errorf("route_id mismatch")
	case ttl <= 0:
		return RouteCooldownState{}, fmt.Errorf("ttl must be positive")
	}

	if state.Until == "" {
		state.Until = now.UTC().Add(ttl).Format(time.RFC3339)
	}
	if state.ExpiresAt == "" {
		state.ExpiresAt = now.UTC().Add(ttl).Format(time.RFC3339)
	}
	return state, nil
}