vibe-proxy/backend/sdk/cliproxy/auth/home_selection.go

package auth

import (
	"context"
	"errors"
	"fmt"
	"slices"
	"strings"
	"sync"
	"sync/atomic"

	"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry"
)

type executionResources struct {
	mu      sync.Mutex
	closed  bool
	closers []func() error
}

type attemptCancel struct {
	cancel context.CancelFunc
	once   sync.Once
}

func (a *attemptCancel) Cancel() {
	if a == nil || a.cancel == nil {
		return
	}
	a.once.Do(a.cancel)
}

type attemptCancels struct {
	mu      sync.Mutex
	closed  bool
	next    uint64
	cancels map[uint64]*attemptCancel
}

func (a *attemptCancels) Add(cancel context.CancelFunc) (func(), error) {
	if a == nil || cancel == nil {
		return func() {}, executionregistry.ErrInvalidExecutionResource
	}

	a.mu.Lock()
	if a.closed {
		a.mu.Unlock()
		cancel()
		return func() {}, executionregistry.ErrRegistryNotAccepting
	}
	if a.cancels == nil {
		a.cancels = make(map[uint64]*attemptCancel)
	}
	a.next++
	token := a.next
	attempt := &attemptCancel{cancel: cancel}
	a.cancels[token] = attempt
	a.mu.Unlock()

	var once sync.Once
	return func() {
		once.Do(func() {
			a.mu.Lock()
			delete(a.cancels, token)
			a.mu.Unlock()
			attempt.Cancel()
		})
	}, nil
}

func (a *attemptCancels) Close() error {
	if a == nil {
		return nil
	}

	a.mu.Lock()
	if a.closed {
		a.mu.Unlock()
		return nil
	}
	a.closed = true
	cancels := a.cancels
	a.cancels = nil
	a.mu.Unlock()

	for _, cancel := range cancels {
		cancel.Cancel()
	}
	return nil
}

func (a *attemptCancels) Len() int {
	if a == nil {
		return 0
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	return len(a.cancels)
}

func (r *executionResources) Add(closeFn func() error) error {
	if closeFn == nil {
		return executionregistry.ErrInvalidExecutionResource
	}

	r.mu.Lock()
	if !r.closed {
		r.closers = append(r.closers, closeFn)
		r.mu.Unlock()
		return nil
	}
	r.mu.Unlock()

	if errClose := closeFn(); errClose != nil {
		return errors.Join(executionregistry.ErrRegistryNotAccepting, errClose)
	}
	return executionregistry.ErrRegistryNotAccepting
}

func (r *executionResources) Close() error {
	r.mu.Lock()
	if r.closed {
		r.mu.Unlock()
		return nil
	}
	r.closed = true
	closers := slices.Clone(r.closers)
	r.closers = nil
	r.mu.Unlock()

	var result error
	for index := len(closers) - 1; index >= 0; index-- {
		result = errors.Join(result, closers[index]())
	}
	return result
}

// HomeDispatchSelection keeps a Home execution scope separate from its auth.
type HomeDispatchSelection struct {
	Auth     *Auth
	Executor ProviderExecutor
	Provider string

	authMu           sync.RWMutex
	scope            *executionregistry.Scope
	accountedModel   string
	requestRetry     int
	hasRequestRetry  bool
	resources        *executionResources
	attemptCancels   *attemptCancels
	once             sync.Once
	retained         atomic.Bool
	runtimeAuthBound atomic.Bool
	ended            atomic.Bool
}

func newHomeDispatchSelection(auth *Auth, executor ProviderExecutor, provider string, scope *executionregistry.Scope) (*HomeDispatchSelection, error) {
	if scope == nil {
		return nil, fmt.Errorf("Home dispatch selection has no execution scope")
	}

	resources := &executionResources{}
	attemptCancels := &attemptCancels{}
	if errBind := resources.Add(attemptCancels.Close); errBind != nil {
		_ = attemptCancels.Close()
		scope.End("attempt_cancel_bind_failed")
		return nil, errBind
	}
	if errBind := scope.Bind(resources.Close); errBind != nil {
		_ = resources.Close()
		scope.End("resource_controller_bind_failed")
		return nil, errBind
	}

	return &HomeDispatchSelection{
		Auth:           auth,
		Executor:       executor,
		Provider:       strings.TrimSpace(provider),
		scope:          scope,
		resources:      resources,
		attemptCancels: attemptCancels,
	}, nil
}

// Bind adds a resource to be closed when this selection ends or drains.
func (s *HomeDispatchSelection) Bind(closeFn func() error) error {
	if s == nil || s.resources == nil {
		if closeFn != nil {
			_ = closeFn()
		}
		return fmt.Errorf("Home dispatch selection has no execution resources")
	}
	return s.resources.Add(closeFn)
}

// AttemptContext creates a selection-owned context and returns its release function.
func (s *HomeDispatchSelection) AttemptContext(ctx context.Context) (context.Context, func(), error) {
	if ctx == nil {
		ctx = context.Background()
	}
	attemptCtx, cancelAttempt := context.WithCancel(ctx)
	if s == nil || s.attemptCancels == nil {
		cancelAttempt()
		return nil, func() {}, fmt.Errorf("Home dispatch selection has no attempt cancels")
	}
	release, errAdd := s.attemptCancels.Add(cancelAttempt)
	if errAdd != nil {
		cancelAttempt()
		return nil, func() {}, errAdd
	}
	return attemptCtx, release, nil
}

// Retain transfers selection ownership from a request to an execution session.
func (s *HomeDispatchSelection) Retain() {
	if s == nil || s.ended.Load() {
		return
	}
	s.retained.Store(true)
}

// Retained reports whether an executor transferred this selection to a session.
func (s *HomeDispatchSelection) Retained() bool {
	return s != nil && s.retained.Load() && !s.ended.Load()
}

// Active reports whether the selection has not ended.
func (s *HomeDispatchSelection) Active() bool {
	return s != nil && !s.ended.Load()
}

// End closes all bound resources and releases the Home execution scope once.
func (s *HomeDispatchSelection) End(reason string) {
	_ = s.EndWithRelease(reason)
}

// EndWithRelease closes all bound resources and returns the Home release ticket.
func (s *HomeDispatchSelection) EndWithRelease(reason string) *executionregistry.ReleaseTicket {
	if s == nil {
		return nil
	}
	var ticket *executionregistry.ReleaseTicket
	s.once.Do(func() {
		s.ended.Store(true)
		if s.scope != nil {
			ticket = s.scope.EndWithRelease(strings.TrimSpace(reason))
		}
	})
	if ticket != nil || s.scope == nil {
		return ticket
	}
	return s.scope.EndWithRelease("")
}

// ReplaceAuth updates the selection after Home returns refreshed credentials.
func (s *HomeDispatchSelection) ReplaceAuth(auth *Auth) {
	if s == nil || auth == nil {
		return
	}
	updated := auth.Clone()
	s.authMu.Lock()
	defer s.authMu.Unlock()
	preserveHomeRoutingAttributes(updated, s.Auth)
	s.Auth = updated
}

func preserveHomeRoutingAttributes(updated, previous *Auth) {
	if updated == nil || previous == nil {
		return
	}
	if updated.Attributes == nil {
		updated.Attributes = make(map[string]string)
	}
	for _, key := range []string{homeUpstreamModelAttributeKey, homeForceMappingAttributeKey, homeOriginalAliasAttributeKey} {
		if value := strings.TrimSpace(previous.Attributes[key]); value != "" {
			updated.Attributes[key] = value
		}
	}
}

// CloneAuth returns a standalone auth copy without the selection handle.
func (s *HomeDispatchSelection) CloneAuth() *Auth {
	if s == nil {
		return nil
	}
	s.authMu.RLock()
	defer s.authMu.RUnlock()
	if s.Auth == nil {
		return nil
	}
	return s.Auth.Clone()
}

// CloneAuthForRoute returns an auth copy adapted for a retained canonical route.
func (s *HomeDispatchSelection) CloneAuthForRoute(routeModel string) *Auth {
	auth := s.CloneAuth()
	if auth == nil || !s.Retained() {
		return auth
	}
	return cloneRetainedHomeAuthForRoute(auth, routeModel)
}

func cloneRetainedHomeAuthForRoute(auth *Auth, routeModel string) *Auth {
	if auth == nil || auth.Attributes == nil {
		return auth
	}
	upstreamModel := strings.TrimSpace(auth.Attributes[homeUpstreamModelAttributeKey])
	if upstreamModel == "" {
		return auth
	}
	upstreamBase, _ := splitRecognizedHomeReasoningSuffix(upstreamModel)
	_, routeSuffix := splitRecognizedHomeReasoningSuffix(routeModel)
	auth.Attributes[homeUpstreamModelAttributeKey] = upstreamBase + routeSuffix
	if strings.EqualFold(strings.TrimSpace(auth.Attributes[homeForceMappingAttributeKey]), "true") {
		auth.Attributes[homeOriginalAliasAttributeKey] = strings.TrimSpace(rewriteModelForAuth(routeModel, auth))
	}
	return auth
}

func splitRecognizedHomeReasoningSuffix(model string) (string, string) {
	model = strings.Trim(model, asciiWhitespace)
	if !strings.HasSuffix(model, ")") {
		return model, ""
	}
	open := strings.LastIndexByte(model, '(')
	if open < 0 || !recognizedHomeConcurrencySuffix(model[open+1:len(model)-1]) {
		return model, ""
	}
	base := strings.Trim(model[:open], asciiWhitespace)
	if base == "" {
		return model, ""
	}
	return base, model[open:]
}