hunjing/shenshi/internal/route/registry.go

// Package route 维护逻辑模型/能力槽位到候选端点的路由注册表。
// 端点来源：环境 SHENSHI_ENDPOINTS（JSON 数组）；未配置时回落 stub（骨架形态）。
// 冷却粒度（shenshi-gateway.md §4）：401/403 只冷却凭证；429/5xx/超时只冷却该端点绑定。
package route

import (
	"encoding/json"
	"fmt"
	"os"
	"sort"
	"time"

	"hunjing/shenshi/internal/transform"
)

// Endpoint 一个可路由的出站端点。
type Endpoint struct {
	Name     string
	Type     string // openai_compatible / local / stub（契约词典 §14.2）
	Adaptor  transform.Adaptor
	Priority int
	// 凭证指纹（凭证级冷却键）；stub/openai 兼容端点各自填充
	CredentialKey string
	// 绑定级冷却键的模型名部分
	UpstreamModel string
}

// EndpointSpec 环境配置形态。
type EndpointSpec struct {
	Type     string `json:"type"`
	Name     string `json:"name"`
	BaseURL  string `json:"base_url"`
	APIKey   string `json:"api_key"`
	Model    string `json:"model"`
	Priority int    `json:"priority"`
}

// Registry 槽位 → 候选端点列表（按 Priority 升序尝试）。
type Registry struct {
	bySlot map[string][]Endpoint
	// Cooldowns/FirstByteTimeoutMS 由 server 层共享使用
	Cooldowns        *Cooldowns
	Sticky           *Sticky
	FirstByteTimeout time.Duration
}

// New 构建注册表：读 SHENSHI_ENDPOINTS（JSON），无配置时登记 stub。
func New() *Registry {
	r := &Registry{
		bySlot:           make(map[string][]Endpoint),
		Cooldowns:        NewCooldowns(),
		Sticky:           NewSticky(),
		FirstByteTimeout: 10 * time.Second,
	}
	specs := loadSpecs()
	if len(specs) == 0 {
		r.Register("chat", Endpoint{Name: "stub", Type: "stub", Adaptor: transform.StubAdaptor{}, Priority: 100, UpstreamModel: "stub"})
		return r
	}
	for _, spec := range specs {
		adaptor, ok := BuildAdaptor(spec)
		if !ok {
			continue
		}
		cred := spec.Name
		if spec.APIKey != "" {
			cred = spec.Name + ":" + fmt.Sprintf("%x", len(spec.APIKey)) // 不落密钥原文，仅指纹
		}
		r.Register("chat", Endpoint{
			Name: spec.Name, Type: spec.Type, Adaptor: adaptor,
			Priority: spec.Priority, CredentialKey: cred, UpstreamModel: spec.Model,
		})
	}
	return r
}

// BuildAdaptor 按端点类型构造出站适配器（契约词典 §14.2；BYOK 动态候选复用）。
func BuildAdaptor(spec EndpointSpec) (transform.Adaptor, bool) {
	switch spec.Type {
	case "openai_compatible", "local":
		return &transform.OpenAICompatAdaptor{
			EndpointName: spec.Name, BaseURL: spec.BaseURL,
			APIKey: spec.APIKey, UpstreamModel: spec.Model,
		}, true
	case "anthropic":
		return &transform.AnthropicAdaptor{
			EndpointName: spec.Name, BaseURL: spec.BaseURL,
			APIKey: spec.APIKey, UpstreamModel: spec.Model,
		}, true
	case "gemini":
		return &transform.GeminiAdaptor{
			EndpointName: spec.Name, BaseURL: spec.BaseURL,
			APIKey: spec.APIKey, UpstreamModel: spec.Model,
		}, true
	default:
		return nil, false
	}
}

func loadSpecs() []EndpointSpec {
	raw := os.Getenv("SHENSHI_ENDPOINTS")
	if raw == "" {
		return nil
	}
	var specs []EndpointSpec
	if err := json.Unmarshal([]byte(raw), &specs); err != nil {
		return nil
	}
	return specs
}

// Reset 清空槽位（测试装配用）。
func (r *Registry) Reset(slot string) { delete(r.bySlot, slot) }

// Register 登记端点到指定能力槽位（契约词典 §14.1）。
func (r *Registry) Register(slot string, ep Endpoint) {
	r.bySlot[slot] = append(r.bySlot[slot], ep)
}

// Candidates 返回槽位候选（Priority 升序；粘性端点置首；冷却中端点排后不剔除——
// 全部冷却时仍可尝试最优先者，避免死锁）。
func (r *Registry) Candidates(slot, runID string) ([]Endpoint, error) {
	if slot == "" {
		slot = "chat"
	}
	eps := append([]Endpoint(nil), r.bySlot[slot]...)
	if len(eps) == 0 {
		return nil, fmt.Errorf("槽位 %s 无可用端点", slot)
	}
	sort.SliceStable(eps, func(i, j int) bool { return eps[i].Priority < eps[j].Priority })
	if runID != "" {
		if stickyName := r.Sticky.Recall(runID); stickyName != "" {
			for i, ep := range eps {
				if ep.Name == stickyName {
					eps = append([]Endpoint{ep}, append(eps[:i:i], eps[i+1:]...)...)
					break
				}
			}
		}
	}
	// 冷却中的端点移到队尾（保留兜底能力）
	cooling := 0
	for i := 0; i < len(eps)-cooling; {
		if r.Cooldowns.Active(bindingKey(eps[i])) || r.Cooldowns.Active(credKey(eps[i])) {
			ep := eps[i]
			eps = append(append(eps[:i:i], eps[i+1:]...), ep)
			cooling++
			continue
		}
		i++
	}
	return eps, nil
}

func bindingKey(ep Endpoint) string {
	return "bind:" + ep.Name + ":" + ep.UpstreamModel
}

func credKey(ep Endpoint) string {
	if ep.CredentialKey == "" {
		return "cred:" + ep.Name
	}
	return "cred:" + ep.CredentialKey
}

// BindingCooldownKey / CredentialCooldownKey 供 server 层标记冷却。
func BindingCooldownKey(ep Endpoint) string    { return bindingKey(ep) }
func CredentialCooldownKey(ep Endpoint) string { return credKey(ep) }