chorus/internal/core/router/routing.go

// Package router contains protocol-independent active-route selection rules.
package router

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"git.ilapage.cn/OPC/chorus/internal/core/model"
	"git.ilapage.cn/OPC/chorus/internal/core/provider"
)

type ErrorCode string

const (
	CodeRouteNotConfigured ErrorCode = "route_not_configured"
	CodeRouteUnavailable   ErrorCode = "route_unavailable"
	CodeFailoverExhausted  ErrorCode = "failover_exhausted"
)

// Error is deliberately code-only so route failures never disclose provider
// configuration, credentials, or internal runtime observations.
type Error struct{ Code ErrorCode }

func (e Error) Error() string { return string(e.Code) }

var (
	ErrRouteNotConfigured = Error{Code: CodeRouteNotConfigured}
	ErrRouteUnavailable   = Error{Code: CodeRouteUnavailable}
	ErrFailoverExhausted  = Error{Code: CodeFailoverExhausted}
	ErrMemberUnavailable  = errors.New("route member is unavailable")
	ErrInvalidSnapshot    = errors.New("route snapshot is invalid")
	ErrInvalidRandom      = errors.New("route random source is invalid")
)

// RouteSnapshot is persisted with a generation at synchronous submission time.
// It carries immutable routing intent but never a Provider URL or credential.
type RouteSnapshot struct {
	Capability            model.Capability `json:"capability"`
	RoutePoolID           uint64           `json:"route_pool_id"`
	RoutePoolVersion      uint32           `json:"route_pool_version"`
	PromptTemplateID      uint64           `json:"prompt_template_id"`
	PromptTemplateKey     string           `json:"prompt_template_key"`
	PromptTemplateVersion uint32           `json:"prompt_template_version"`
	MaxFailover           uint16           `json:"max_failover"`
	Members               []MemberSnapshot `json:"members"`
}

type MemberSnapshot struct {
	RoutePoolMemberID uint64 `json:"route_pool_member_id"`
	ProviderModelID   uint64 `json:"provider_model_id"`
	Weight            uint16 `json:"weight"`
	FailureThreshold  uint32 `json:"failure_threshold"`
	OpenSeconds       uint32 `json:"open_seconds"`
	HalfOpenMax       uint16 `json:"half_open_max"`
}

func (s RouteSnapshot) Validate() error {
	if !s.Capability.Valid() || s.RoutePoolID == 0 || s.RoutePoolVersion == 0 || s.PromptTemplateID == 0 || s.PromptTemplateVersion == 0 || s.PromptTemplateKey == "" || len(s.Members) == 0 {
		return ErrInvalidSnapshot
	}
	members := make(map[uint64]struct{}, len(s.Members))
	models := make(map[uint64]struct{}, len(s.Members))
	for _, member := range s.Members {
		if member.RoutePoolMemberID == 0 || member.ProviderModelID == 0 || member.Weight == 0 || member.FailureThreshold == 0 || member.OpenSeconds == 0 || member.HalfOpenMax == 0 {
			return ErrInvalidSnapshot
		}
		if _, exists := members[member.RoutePoolMemberID]; exists {
			return ErrInvalidSnapshot
		}
		if _, exists := models[member.ProviderModelID]; exists {
			return ErrInvalidSnapshot
		}
		members[member.RoutePoolMemberID] = struct{}{}
		models[member.ProviderModelID] = struct{}{}
	}
	return nil
}

// Encode validates and serializes the exact non-secret snapshot saved to a
// generation. Provider credentials remain database-only and are loaded later.
func (s RouteSnapshot) Encode() (json.RawMessage, error) {
	if err := s.Validate(); err != nil {
		return nil, err
	}
	encoded, err := json.Marshal(s)
	if err != nil {
		return nil, fmt.Errorf("encode route snapshot: %w", err)
	}
	return json.RawMessage(encoded), nil
}

// ApplySnapshot prepares a pending generation for persistence without calling
// an upstream service. Idempotent creation still owns the database transaction.
func ApplySnapshot(generation *model.Generation, snapshot RouteSnapshot) error {
	if generation == nil {
		return ErrInvalidSnapshot
	}
	encoded, err := snapshot.Encode()
	if err != nil {
		return err
	}
	routePoolID, routePoolVersion := snapshot.RoutePoolID, snapshot.RoutePoolVersion
	promptTemplateID := snapshot.PromptTemplateID
	generation.RoutePoolID = &routePoolID
	generation.RoutePoolVersion = &routePoolVersion
	generation.PromptTemplateID = &promptTemplateID
	generation.RouteSnapshot = append(generation.RouteSnapshot[:0], encoded...)
	return nil
}

// CapabilityForSubmission maps existing generation input shape to the MVP-1
// capability contract. Future UI modes can call this without protocol guesses.
func CapabilityForSubmission(kind model.GenerationKind, hasInputs bool) (model.Capability, error) {
	switch kind {
	case model.KindText:
		return model.CapabilityText, nil
	case model.KindImage:
		if hasInputs {
			return model.CapabilityImageEdit, nil
		}
		return model.CapabilityImageGenerate, nil
	default:
		return "", ErrRouteNotConfigured
	}
}

type CircuitState string

const (
	CircuitClosed   CircuitState = "closed"
	CircuitOpen     CircuitState = "open"
	CircuitHalfOpen CircuitState = "half_open"
)

// MemberState is the current, mutable eligibility view. Snapshot members are
// rechecked against it before worker selection so disabled configuration wins.
type MemberState struct {
	RoutePoolMemberID  uint64
	Enabled            bool
	ProviderEnabled    bool
	ModelEnabled       bool
	SupportsCapability bool
	CircuitState       CircuitState
	HalfOpenInFlight   uint16
	HalfOpenMax        uint16
}

func (s MemberState) Eligible() bool {
	if !s.Enabled || !s.ProviderEnabled || !s.ModelEnabled || !s.SupportsCapability {
		return false
	}
	switch s.CircuitState {
	case CircuitClosed:
		return true
	case CircuitHalfOpen:
		return s.HalfOpenMax > 0 && s.HalfOpenInFlight < s.HalfOpenMax
	default:
		return false
	}
}

// RandomSource is injected so weighted selection is deterministic in tests.
type RandomSource interface{ Uint64n(uint64) uint64 }

// SelectCandidates performs weighted sampling without replacement. At most the
// initial attempt plus MaxFailover alternatives are returned.
func SelectCandidates(snapshot RouteSnapshot, states []MemberState, random RandomSource) ([]MemberSnapshot, error) {
	if err := snapshot.Validate(); err != nil {
		return nil, err
	}
	if random == nil {
		return nil, ErrInvalidRandom
	}
	stateByMember := make(map[uint64]MemberState, len(states))
	for _, state := range states {
		stateByMember[state.RoutePoolMemberID] = state
	}
	available := make([]MemberSnapshot, 0, len(snapshot.Members))
	for _, member := range snapshot.Members {
		state, ok := stateByMember[member.RoutePoolMemberID]
		if ok && state.Eligible() {
			available = append(available, member)
		}
	}
	if len(available) == 0 {
		return nil, ErrRouteUnavailable
	}
	limit := int(snapshot.MaxFailover) + 1
	if limit > len(available) {
		limit = len(available)
	}
	selected := make([]MemberSnapshot, 0, limit)
	for len(selected) < limit {
		var total uint64
		for _, member := range available {
			total += uint64(member.Weight)
		}
		pick := random.Uint64n(total)
		if pick >= total {
			return nil, ErrInvalidRandom
		}
		var cumulative uint64
		chosen := 0
		for index, member := range available {
			cumulative += uint64(member.Weight)
			if pick < cumulative {
				chosen = index
				break
			}
		}
		selected = append(selected, available[chosen])
		available = append(available[:chosen], available[chosen+1:]...)
	}
	return selected, nil
}

// CandidateForAttempt returns the next no-repeat candidate after a retryable
// attempt. Callers must stop immediately for non-retryable provider failures.
func CandidateForAttempt(candidates []MemberSnapshot, providerAttemptCount int) (MemberSnapshot, error) {
	if providerAttemptCount < 0 {
		return MemberSnapshot{}, ErrInvalidSnapshot
	}
	if len(candidates) == 0 {
		return MemberSnapshot{}, ErrRouteUnavailable
	}
	if providerAttemptCount >= len(candidates) {
		return MemberSnapshot{}, ErrFailoverExhausted
	}
	return candidates[providerAttemptCount], nil
}

type FailureAction string

const (
	FailureTryNext FailureAction = "try_next"
	FailureStop    FailureAction = "stop"
)

func ActionForFailure(class provider.FailureClass) FailureAction {
	if Retryable(class) {
		return FailureTryNext
	}
	return FailureStop
}

// SnapshotRepository is the read boundary used by synchronous submission. Its
// implementation must return ErrRouteNotConfigured for no active binding and
// ErrRouteUnavailable for an active binding with no usable members.
type SnapshotRepository interface {
	ActiveSnapshot(context.Context, model.Capability) (RouteSnapshot, error)
}

// RuntimeRepository is the database-atomic boundary consumed by the worker in
// #22. The owner/token lease lets implementations cap concurrent half-open
// probes without weakening the circuit across worker processes.
type RuntimeRepository interface {
	MemberStates(context.Context, model.Capability, []MemberSnapshot) ([]MemberState, error)
	Reserve(context.Context, ReservationRequest) (Reservation, error)
	Record(context.Context, Reservation, CircuitObservation) (bool, error)
}

type ReservationRequest struct {
	RoutePoolMemberID uint64
	Capability        model.Capability
	Owner             string
	LeaseDuration     time.Duration
}

type Reservation struct {
	RoutePoolMemberID uint64
	Token             string
	State             CircuitState
}

type CircuitObservation struct {
	Retryable       bool
	Succeeded       bool
	OpenImmediately bool
	ErrorCode       string
	ErrorMessage    string
}