easyai-ai-gateway/apps/api/internal/modelaccess/scopes.go

package modelaccess

import "strings"

// ScopeAllowsTask reports whether scopes authorize one public task kind.
// Empty scopes retain the legacy unrestricted behavior for old API keys.
func ScopeAllowsTask(scopes []string, kind string) bool {
	if len(scopes) == 0 {
		return true
	}
	return scopeAllowsCapability(scopes, capabilityForTaskKind(kind))
}

// ScopeAllowsModelType reports whether scopes authorize one runtime model type.
// Unknown model types are deny-by-default unless the key has all or an exact
// custom scope matching that model type.
func ScopeAllowsModelType(scopes []string, modelType string) bool {
	if len(scopes) == 0 {
		return true
	}
	modelType = normalize(modelType)
	if modelType == "" {
		return false
	}
	capability := capabilityForModelType(modelType)
	if capability == "" {
		return hasScope(scopes, modelType)
	}
	return scopeAllowsCapability(scopes, capability)
}

// FilterModelTypes keeps the declared order while removing model types that
// the key cannot invoke.
func FilterModelTypes(scopes []string, modelTypes []string) []string {
	if len(modelTypes) == 0 {
		return nil
	}
	filtered := make([]string, 0, len(modelTypes))
	seen := map[string]bool{}
	for _, modelType := range modelTypes {
		normalized := normalize(modelType)
		if normalized == "" || seen[normalized] || !ScopeAllowsModelType(scopes, normalized) {
			continue
		}
		seen[normalized] = true
		filtered = append(filtered, normalized)
	}
	return filtered
}

func capabilityForTaskKind(kind string) string {
	switch normalize(kind) {
	case "chat.completions", "responses":
		return "chat"
	case "embeddings":
		return "embedding"
	case "reranks":
		return "rerank"
	case "images.generations", "images.edits":
		return "image"
	case "images.vectorize":
		return "image_vectorize"
	case "videos.generations":
		return "video"
	case "videos.upscales":
		return "video_enhance"
	case "song.generations", "music.generations":
		return "music"
	case "speech.generations":
		return "audio"
	case "voice.clone":
		return "voice_clone"
	default:
		return normalize(kind)
	}
}

func capabilityForModelType(modelType string) string {
	switch normalize(modelType) {
	case "text_generate", "tools_call":
		return "chat"
	case "text_embedding":
		return "embedding"
	case "text_rerank":
		return "rerank"
	case "image_generate", "image_edit", "image_analysis":
		return "image"
	case "image_vectorize":
		return "image_vectorize"
	case "video_generate", "image_to_video", "text_to_video", "video_edit", "video_reference", "video_first_last_frame", "video_understanding", "omni_video", "omni":
		return "video"
	case "video_enhance":
		return "video_enhance"
	case "audio_generate", "music_generate":
		return "music"
	case "text_to_speech", "audio_understanding":
		return "audio"
	case "voice_clone":
		return "voice_clone"
	default:
		return ""
	}
}

func scopeAllowsCapability(scopes []string, capability string) bool {
	capability = normalize(capability)
	if capability == "" {
		return false
	}
	for _, scope := range scopes {
		scope = normalize(scope)
		if scope == "*" || scope == "all" || scope == capability {
			return true
		}
		switch capability {
		case "chat":
			if scope == "text" || scope == "text_generate" {
				return true
			}
		case "embedding":
			if scope == "text_embedding" {
				return true
			}
		case "rerank":
			if scope == "text_rerank" {
				return true
			}
		case "music":
			if scope == "audio_generate" || scope == "music_generate" || scope == "song" {
				return true
			}
		case "audio":
			if scope == "text_to_speech" || scope == "speech" || scope == "tts" {
				return true
			}
		case "voice_clone":
			if scope == "audio" || scope == "text_to_speech" || scope == "speech" || scope == "tts" {
				return true
			}
		case "image_vectorize":
			if scope == "image" || scope == "vectorize" {
				return true
			}
		case "video_enhance":
			if scope == "video" || scope == "video_upscale" || scope == "upscale" {
				return true
			}
		}
	}
	return false
}

func hasScope(scopes []string, want string) bool {
	want = normalize(want)
	for _, scope := range scopes {
		scope = normalize(scope)
		if scope == "*" || scope == "all" || scope == want {
			return true
		}
	}
	return false
}

func normalize(value string) string {
	return strings.ReplaceAll(strings.ToLower(strings.TrimSpace(value)), "-", "_")
}