iop/apps/edge/internal/openai/principal_routes.go

package openai

import (
	"context"
	"errors"
	"net/http"
	"sort"
	"strings"

	"iop/apps/edge/internal/authprojection"
	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/config"
)

var (
	ErrRouteNotFound     = errors.New("route not found")
	ErrProjectionExpired = errors.New("projection is expired")
	ErrPrincipalRequired = errors.New("principal is required")
)

func (s *Server) advertisedModelsForPrincipal(ctx context.Context) ([]advertisedModel, error) {
	projection := s.PrincipalProjection()
	if projection == nil || projection.State() == authprojection.StateUnmanaged {
		return s.advertisedModels(), nil
	}
	view, ok := authenticatedProjectionViewFromContext(ctx)
	if !ok || strings.TrimSpace(view.Principal.PrincipalRef) == "" {
		if projection.State() == authprojection.StateExpired {
			return nil, ErrProjectionExpired
		}
		return nil, ErrPrincipalRequired
	}
	routes := view.Routes

	seen := make(map[string]struct{})
	var ids []string
	for _, r := range routes {
		id := strings.TrimSpace(r.RouteID)
		if id != "" {
			if _, exists := seen[id]; !exists {
				seen[id] = struct{}{}
				ids = append(ids, id)
			}
		}
	}
	sort.Strings(ids)

	models := make([]advertisedModel, 0, len(ids))
	for _, id := range ids {
		models = append(models, advertisedModel{
			ID:          id,
			DisplayName: id,
		})
	}
	return models, nil
}

func (s *Server) resolveRouteDispatchForPrincipal(ctx context.Context, model string) (routeDispatch, error) {
	projection := s.PrincipalProjection()
	if _, ok := authenticatedProjectionViewFromContext(ctx); ok {
		return s.resolveProjectedRoute(ctx, model)
	}
	if projection != nil && projection.State() != authprojection.StateUnmanaged {
		if projection.State() == authprojection.StateExpired {
			return routeDispatch{}, ErrProjectionExpired
		}
		return routeDispatch{}, ErrPrincipalRequired
	}
	dispatch, ok := s.resolveRouteDispatch(model)
	if !ok {
		return routeDispatch{}, ErrRouteNotFound
	}
	return dispatch, nil
}

func (s *Server) resolveProjectedRoute(ctx context.Context, model string) (routeDispatch, error) {
	model = strings.TrimSpace(model)
	if model == "" {
		return routeDispatch{}, ErrRouteNotFound
	}

	p, ok := principalFromContext(ctx)
	if !ok || strings.TrimSpace(p.PrincipalRef) == "" {
		return routeDispatch{}, ErrPrincipalRequired
	}

	view, ok := authenticatedProjectionViewFromContext(ctx)
	if !ok || view.Principal.PrincipalRef != p.PrincipalRef {
		return routeDispatch{}, ErrPrincipalRequired
	}
	routes := view.Routes

	var matchedRoute *authprojection.Route
	for i := range routes {
		r := &routes[i]
		if r.RouteID == model || (r.RouteAlias != "" && r.RouteAlias == model) {
			matchedRoute = r
			break
		}
	}

	if matchedRoute == nil {
		return routeDispatch{}, ErrRouteNotFound
	}

	binding, err := resolveManagedCatalogBinding(*matchedRoute, s.modelCatalogSnapshot())
	if err != nil {
		return routeDispatch{}, err
	}
	pred := managedRouteCandidatePredicate(*matchedRoute, binding.ProviderID)

	return routeDispatch{
		NodeRef:              s.cfg.NodeRef,
		ProviderID:           binding.ProviderID,
		UsageAttribution:     config.UsageAttributionProvider,
		SessionID:            s.resolveSessionID(),
		TimeoutSec:           s.resolveTimeoutSec(),
		ProviderPool:         true,
		Managed:              true,
		ModelGroupKey:        binding.ModelGroupKey,
		RouteID:              matchedRoute.RouteID,
		CredentialSlotRef:    matchedRoute.CredentialSlotRef,
		ProfileID:            matchedRoute.ProfileID,
		UpstreamModel:        matchedRoute.UpstreamModel,
		ResourceSelector:     matchedRoute.ResourceSelector,
		RouteRevision:        matchedRoute.RouteRevision,
		CredentialRevision:   matchedRoute.CredentialRevision,
		PrincipalRef:         matchedRoute.PrincipalRef,
		ProjectionGeneration: view.Generation,
		ManagedPredicate:     pred,
	}, nil
}

type managedCatalogBinding struct{ ModelGroupKey, ProviderID string }

func resolveManagedCatalogBinding(route authprojection.Route, catalog []config.ModelCatalogEntry) (managedCatalogBinding, error) {
	selector := strings.TrimSpace(route.ResourceSelector)
	if selector == "" {
		return managedCatalogBinding{}, ErrRouteNotFound
	}
	explicit := !strings.EqualFold(selector, "default")
	var matches []managedCatalogBinding
	for _, entry := range catalog {
		group := strings.TrimSpace(entry.ID)
		if group == "" {
			continue
		}
		for providerID, servedModel := range entry.Providers {
			if strings.TrimSpace(servedModel) != strings.TrimSpace(route.UpstreamModel) {
				continue
			}
			if explicit && strings.TrimSpace(providerID) != selector {
				continue
			}
			binding := managedCatalogBinding{ModelGroupKey: group}
			if explicit {
				binding.ProviderID = strings.TrimSpace(providerID)
			}
			matches = append(matches, binding)
			break
		}
	}
	if len(matches) != 1 {
		return managedCatalogBinding{}, ErrRouteNotFound
	}
	return matches[0], nil
}

func managedRouteCandidatePredicate(route authprojection.Route, providerID string) edgeservice.ProviderPoolCandidatePredicate {
	return func(c edgeservice.ProviderPoolCandidate) bool {
		if c.ExecutionPath != string(edgeservice.ProviderPoolPathTunnel) {
			return false
		}
		if providerID != "" && c.ProviderID != providerID {
			return false
		}
		if route.ProfileID != "" && c.ProfileID != route.ProfileID {
			return false
		}
		if route.UpstreamModel != "" && c.ActualModel != route.UpstreamModel {
			return false
		}
		return true
	}
}

func composeCandidatePredicates(p1, p2 edgeservice.ProviderPoolCandidatePredicate) edgeservice.ProviderPoolCandidatePredicate {
	if p1 == nil {
		return p2
	}
	if p2 == nil {
		return p1
	}
	return func(c edgeservice.ProviderPoolCandidate) bool {
		return p1(c) && p2(c)
	}
}

func (s *Server) writeManagedRouteError(w http.ResponseWriter, r *http.Request, err error) {
	if isAnthropicRequest(r) {
		s.writeAnthropicRouteError(w, err)
		return
	}
	writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
}

func (s *Server) writeAnthropicRouteError(w http.ResponseWriter, err error) {
	if errors.Is(err, ErrProjectionExpired) {
		writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
		return
	}
	writeAnthropicError(w, http.StatusBadRequest, "not_supported_error", "model does not resolve to a protocol profile")
}