feat: 实现路由核心与故障转移契约 (#21)

This commit is contained in:
ila
2026-08-21 17:55:16 +08:00
parent 09e8601bb8
commit 9b2827a727
7 changed files with 743 additions and 32 deletions
+38 -29
View File
@@ -49,40 +49,48 @@ type ProviderModel struct {
func (ProviderModel) TableName() string { return "provider_models" }
type PromptTemplate struct {
ID uint64 `gorm:"column:id;primaryKey;autoIncrement"`
TemplateKey string `gorm:"column:template_key;size:120;not null;uniqueIndex:uq_prompt_templates_key"`
Kind GenerationKind `gorm:"column:kind;size:16;not null"`
APIType APIType `gorm:"column:api_type;size:32;not null"`
Name string `gorm:"column:name;size:120;not null"`
Version uint32 `gorm:"column:version;not null"`
TemplateText string `gorm:"column:template_text;type:text;not null"`
Enabled bool `gorm:"column:enabled;not null"`
CreatedAt time.Time `gorm:"column:created_at;not null"`
UpdatedAt time.Time `gorm:"column:updated_at;not null"`
ID uint64 `gorm:"column:id;primaryKey;autoIncrement"`
TemplateKey string `gorm:"column:template_key;size:120;not null;uniqueIndex:uq_prompt_templates_key"`
Kind GenerationKind `gorm:"column:kind;size:16;not null"`
APIType APIType `gorm:"column:api_type;size:32;not null"`
Capability Capability `gorm:"column:capability;size:32;not null"`
Name string `gorm:"column:name;size:120;not null"`
Version uint32 `gorm:"column:version;not null"`
TemplateText string `gorm:"column:template_text;type:text;not null"`
DefaultRoleRule string `gorm:"column:default_role_rule;type:text;not null"`
Enabled bool `gorm:"column:enabled;not null"`
CreatedAt time.Time `gorm:"column:created_at;not null"`
UpdatedAt time.Time `gorm:"column:updated_at;not null"`
}
func (PromptTemplate) TableName() string { return "prompt_templates" }
type Generation struct {
ID uint64 `gorm:"column:id;primaryKey;autoIncrement"`
UserID uint64 `gorm:"column:user_id;not null;uniqueIndex:uq_generations_user_idempotency,priority:1"`
ProviderModelID *uint64 `gorm:"column:provider_model_id"`
Kind GenerationKind `gorm:"column:kind;size:16;not null"`
Status GenerationStatus `gorm:"column:status;size:16;not null"`
IdempotencyKey string `gorm:"column:idempotency_key;size:128;not null;uniqueIndex:uq_generations_user_idempotency,priority:2"`
UserPrompt string `gorm:"column:user_prompt;type:text;not null"`
RenderedPrompt string `gorm:"column:rendered_prompt;type:text;not null"`
Attempts json.RawMessage `gorm:"column:attempts;type:json;not null"`
AttemptCount uint32 `gorm:"column:attempt_count;not null"`
ErrorCode *string `gorm:"column:error_code;size:64"`
ErrorMessage *string `gorm:"column:error_message;size:1024"`
LeaseOwner *string `gorm:"column:lease_owner;size:128"`
LeaseToken *string `gorm:"column:lease_token;size:36"`
LeaseUntil *time.Time `gorm:"column:lease_until"`
StartedAt *time.Time `gorm:"column:started_at"`
CompletedAt *time.Time `gorm:"column:completed_at"`
CreatedAt time.Time `gorm:"column:created_at;not null"`
UpdatedAt time.Time `gorm:"column:updated_at;not null"`
ID uint64 `gorm:"column:id;primaryKey;autoIncrement"`
UserID uint64 `gorm:"column:user_id;not null;uniqueIndex:uq_generations_user_idempotency,priority:1"`
ProviderModelID *uint64 `gorm:"column:provider_model_id"`
RoutePoolID *uint64 `gorm:"column:route_pool_id"`
RoutePoolVersion *uint32 `gorm:"column:route_pool_version"`
PromptTemplateID *uint64 `gorm:"column:prompt_template_id"`
RouteSnapshot json.RawMessage `gorm:"column:route_snapshot;type:json"`
RoleRule *string `gorm:"column:role_rule;type:text"`
ProviderAttemptCount uint32 `gorm:"column:provider_attempt_count;not null"`
Kind GenerationKind `gorm:"column:kind;size:16;not null"`
Status GenerationStatus `gorm:"column:status;size:16;not null"`
IdempotencyKey string `gorm:"column:idempotency_key;size:128;not null;uniqueIndex:uq_generations_user_idempotency,priority:2"`
UserPrompt string `gorm:"column:user_prompt;type:text;not null"`
RenderedPrompt string `gorm:"column:rendered_prompt;type:text;not null"`
Attempts json.RawMessage `gorm:"column:attempts;type:json;not null"`
AttemptCount uint32 `gorm:"column:attempt_count;not null"`
ErrorCode *string `gorm:"column:error_code;size:64"`
ErrorMessage *string `gorm:"column:error_message;size:1024"`
LeaseOwner *string `gorm:"column:lease_owner;size:128"`
LeaseToken *string `gorm:"column:lease_token;size:36"`
LeaseUntil *time.Time `gorm:"column:lease_until"`
StartedAt *time.Time `gorm:"column:started_at"`
CompletedAt *time.Time `gorm:"column:completed_at"`
CreatedAt time.Time `gorm:"column:created_at;not null"`
UpdatedAt time.Time `gorm:"column:updated_at;not null"`
}
func (Generation) TableName() string { return "generations" }
@@ -92,6 +100,7 @@ type GenerationInput struct {
GenerationID uint64 `gorm:"column:generation_id;not null;uniqueIndex:uq_generation_inputs_position,priority:1"`
Position uint32 `gorm:"column:position;not null;uniqueIndex:uq_generation_inputs_position,priority:2"`
Role InputRole `gorm:"column:role;size:16;not null"`
Note *string `gorm:"column:note;size:500"`
OriginalName string `gorm:"column:original_name;size:255;not null"`
MIMEType string `gorm:"column:mime_type;size:120;not null"`
StorageKey string `gorm:"column:storage_key;size:512;not null"`
+3 -3
View File
@@ -16,9 +16,9 @@ func TestGORMModelsMatchMigrationColumns(t *testing.T) {
{User{}, []string{"id", "email", "password_hash", "display_name", "status", "created_at", "updated_at"}},
{Provider{}, []string{"id", "slug", "name", "base_url", "auth_type", "api_key_enc", "enabled", "created_at", "updated_at"}},
{ProviderModel{}, []string{"id", "provider_id", "name", "model_id", "api_type", "kind", "extra_body", "timeout_ms", "weight", "enabled", "created_at", "updated_at"}},
{PromptTemplate{}, []string{"id", "template_key", "kind", "api_type", "name", "version", "template_text", "enabled", "created_at", "updated_at"}},
{Generation{}, []string{"id", "user_id", "provider_model_id", "kind", "status", "idempotency_key", "user_prompt", "rendered_prompt", "attempts", "attempt_count", "error_code", "error_message", "lease_owner", "lease_token", "lease_until", "started_at", "completed_at", "created_at", "updated_at"}},
{GenerationInput{}, []string{"id", "generation_id", "position", "role", "original_name", "mime_type", "storage_key", "size_bytes", "width", "height", "created_at"}},
{PromptTemplate{}, []string{"id", "template_key", "kind", "api_type", "capability", "name", "version", "template_text", "default_role_rule", "enabled", "created_at", "updated_at"}},
{Generation{}, []string{"id", "user_id", "provider_model_id", "route_pool_id", "route_pool_version", "prompt_template_id", "route_snapshot", "role_rule", "provider_attempt_count", "kind", "status", "idempotency_key", "user_prompt", "rendered_prompt", "attempts", "attempt_count", "error_code", "error_message", "lease_owner", "lease_token", "lease_until", "started_at", "completed_at", "created_at", "updated_at"}},
{GenerationInput{}, []string{"id", "generation_id", "position", "role", "note", "original_name", "mime_type", "storage_key", "size_bytes", "width", "height", "created_at"}},
{GenerationOutput{}, []string{"id", "generation_id", "kind", "text_content", "storage_key", "thumbnail_storage_key", "mime_type", "size_bytes", "width", "height", "created_at"}},
}
+14
View File
@@ -74,3 +74,17 @@ const (
APIImagesEdits APIType = "images_edits"
APIGemini APIType = "gemini"
)
// Capability is the product-level operation selected by an active route.
// It intentionally does not infer support from a provider protocol.
type Capability string
const (
CapabilityText Capability = "text"
CapabilityImageGenerate Capability = "image_generate"
CapabilityImageEdit Capability = "image_edit"
)
func (c Capability) Valid() bool {
return c == CapabilityText || c == CapabilityImageGenerate || c == CapabilityImageEdit
}
+109
View File
@@ -0,0 +1,109 @@
package router
import (
"context"
"errors"
"fmt"
"git.ilapage.cn/OPC/chorus/internal/core/model"
"gorm.io/gorm"
)
// GORMRepository resolves an active route into the no-secret snapshot stored
// on a generation. It deliberately never selects Provider base URLs or keys.
type GORMRepository struct{ db *gorm.DB }
func NewGORMRepository(db *gorm.DB) (*GORMRepository, error) {
if db == nil {
return nil, errors.New("route database is required")
}
return &GORMRepository{db: db}, nil
}
func (r *GORMRepository) ActiveSnapshot(ctx context.Context, capability model.Capability) (RouteSnapshot, error) {
if !capability.Valid() {
return RouteSnapshot{}, ErrRouteNotConfigured
}
type routeRow struct {
RoutePoolID, PromptTemplateID uint64
RoutePoolVersion, PromptTemplateVersion uint32
PromptTemplateKey string
MaxFailover uint16
}
var route routeRow
result := r.db.WithContext(ctx).Table("active_routes ar").
Select(`rp.id AS route_pool_id, rp.version AS route_pool_version,
pt.id AS prompt_template_id, pt.template_key AS prompt_template_key,
pt.version AS prompt_template_version, rp.max_failover`).
Joins("JOIN route_pools rp ON rp.id = ar.route_pool_id").
Joins("JOIN prompt_templates pt ON pt.id = rp.prompt_template_id").
Where("ar.capability = ? AND rp.enabled = TRUE AND rp.capability = ? AND pt.enabled = TRUE AND pt.capability = ?", capability, capability, capability).
Limit(1).Scan(&route)
if result.Error != nil {
return RouteSnapshot{}, fmt.Errorf("load active route: %w", result.Error)
}
if result.RowsAffected == 0 {
return RouteSnapshot{}, r.routeAvailabilityError(ctx, capability)
}
type memberRow struct {
RoutePoolMemberID, ProviderModelID uint64
Weight uint16
FailureThreshold, OpenSeconds uint32
HalfOpenMax uint16
}
var rows []memberRow
result = r.db.WithContext(ctx).Table("route_pool_members rpm").
Select(`rpm.id AS route_pool_member_id, rpm.provider_model_id,
rpm.weight, rpm.failure_threshold, rpm.open_seconds, rpm.half_open_max`).
Joins("JOIN provider_models pm ON pm.id = rpm.provider_model_id").
Joins("JOIN providers p ON p.id = pm.provider_id").
Joins("JOIN provider_model_capabilities pmc ON pmc.provider_model_id = pm.id AND pmc.capability = ?", capability).
Where("rpm.route_pool_id = ? AND rpm.enabled = TRUE AND pm.enabled = TRUE AND p.enabled = TRUE", route.RoutePoolID).
Order("rpm.position, rpm.id").Scan(&rows)
if result.Error != nil {
return RouteSnapshot{}, fmt.Errorf("load route members: %w", result.Error)
}
if len(rows) == 0 {
return RouteSnapshot{}, ErrRouteUnavailable
}
snapshot := RouteSnapshot{
Capability: capability,
RoutePoolID: route.RoutePoolID,
RoutePoolVersion: route.RoutePoolVersion,
PromptTemplateID: route.PromptTemplateID,
PromptTemplateKey: route.PromptTemplateKey,
PromptTemplateVersion: route.PromptTemplateVersion,
MaxFailover: route.MaxFailover,
Members: make([]MemberSnapshot, 0, len(rows)),
}
for _, row := range rows {
snapshot.Members = append(snapshot.Members, MemberSnapshot{
RoutePoolMemberID: row.RoutePoolMemberID,
ProviderModelID: row.ProviderModelID,
Weight: row.Weight,
FailureThreshold: row.FailureThreshold,
OpenSeconds: row.OpenSeconds,
HalfOpenMax: row.HalfOpenMax,
})
}
if err := snapshot.Validate(); err != nil {
return RouteSnapshot{}, ErrRouteUnavailable
}
return snapshot, nil
}
func (r *GORMRepository) routeAvailabilityError(ctx context.Context, capability model.Capability) error {
var count int64
result := r.db.WithContext(ctx).Table("active_routes").Where("capability = ?", capability).Count(&count)
if result.Error != nil {
return fmt.Errorf("check active route: %w", result.Error)
}
if count == 0 {
return ErrRouteNotConfigured
}
return ErrRouteUnavailable
}
var _ SnapshotRepository = (*GORMRepository)(nil)
@@ -0,0 +1,116 @@
package router
import (
"context"
"encoding/json"
"fmt"
"os"
"testing"
"time"
"git.ilapage.cn/OPC/chorus/internal/core/model"
"gorm.io/driver/mysql"
"gorm.io/gorm"
"gorm.io/gorm/logger"
)
func TestGORMRepositoryActiveSnapshotMySQL(t *testing.T) {
dsn := os.Getenv("CHORUS_TEST_DSN")
if dsn == "" {
t.Skip("CHORUS_TEST_DSN is not set")
}
db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
if err != nil {
t.Fatal(err)
}
sqlDB, err := db.DB()
if err != nil {
t.Fatal(err)
}
defer sqlDB.Close()
tx := db.Begin()
if tx.Error != nil {
t.Fatal(tx.Error)
}
defer tx.Rollback()
repository, err := NewGORMRepository(tx)
if err != nil {
t.Fatal(err)
}
suffix := fmt.Sprintf("router-%d", time.Now().UnixNano())
capability, apiType, kind := availableTestCapability(t, db)
provider := model.Provider{Slug: suffix, Name: "Router Test Provider", BaseURL: "https://provider.invalid/v1", AuthType: "none", Enabled: true}
if err := tx.Create(&provider).Error; err != nil {
t.Fatal(err)
}
providerModel := model.ProviderModel{ProviderID: provider.ID, Name: "Router Test Model", ModelID: suffix, APIType: apiType, Kind: kind, ExtraBody: json.RawMessage("{}"), TimeoutMS: 1000, Weight: 100, Enabled: true}
if err := tx.Create(&providerModel).Error; err != nil {
t.Fatal(err)
}
if err := tx.Exec("INSERT INTO provider_model_capabilities (provider_model_id, capability) VALUES (?, ?)", providerModel.ID, capability).Error; err != nil {
t.Fatal(err)
}
template := model.PromptTemplate{TemplateKey: suffix, Kind: kind, APIType: apiType, Capability: capability, Name: "Router Test Template", Version: 2, TemplateText: "{{.UserPrompt}}", DefaultRoleRule: "", Enabled: true}
if err := tx.Create(&template).Error; err != nil {
t.Fatal(err)
}
result := tx.Exec(`INSERT INTO route_pools (slug, name, capability, prompt_template_id, max_failover, version, enabled) VALUES (?, ?, ?, ?, 1, 4, TRUE)`, suffix, "Router Test Pool", capability, template.ID)
if result.Error != nil {
t.Fatal(result.Error)
}
var routePoolID uint64
if err := tx.Raw("SELECT id FROM route_pools WHERE slug = ?", suffix).Scan(&routePoolID).Error; err != nil {
t.Fatal(err)
}
if err := tx.Exec("INSERT INTO active_routes (capability, route_pool_id) VALUES (?, ?)", capability, routePoolID).Error; err != nil {
t.Fatal(err)
}
if err := tx.Exec(`INSERT INTO route_pool_members (route_pool_id, provider_model_id, weight, failure_threshold, open_seconds, half_open_max, enabled, position) VALUES (?, ?, 5, 2, 45, 2, TRUE, 0)`, routePoolID, providerModel.ID).Error; err != nil {
t.Fatal(err)
}
snapshot, err := repository.ActiveSnapshot(context.Background(), capability)
if err != nil {
t.Fatal(err)
}
if snapshot.RoutePoolID != routePoolID || snapshot.RoutePoolVersion != 4 || snapshot.PromptTemplateID != template.ID || snapshot.PromptTemplateVersion != 2 || snapshot.MaxFailover != 1 || len(snapshot.Members) != 1 || snapshot.Members[0].ProviderModelID != providerModel.ID || snapshot.Members[0].Weight != 5 {
t.Fatalf("snapshot = %#v", snapshot)
}
if err := tx.Exec("UPDATE route_pool_members SET enabled = FALSE WHERE route_pool_id = ?", routePoolID).Error; err != nil {
t.Fatal(err)
}
if _, err := repository.ActiveSnapshot(context.Background(), capability); err != ErrRouteUnavailable {
t.Fatalf("disabled member error = %v", err)
}
if err := tx.Exec("DELETE FROM active_routes WHERE capability = ?", capability).Error; err != nil {
t.Fatal(err)
}
if _, err := repository.ActiveSnapshot(context.Background(), capability); err != ErrRouteNotConfigured {
t.Fatalf("missing active route error = %v", err)
}
}
func availableTestCapability(t *testing.T, db *gorm.DB) (model.Capability, model.APIType, model.GenerationKind) {
t.Helper()
for _, candidate := range []struct {
capability model.Capability
apiType model.APIType
kind model.GenerationKind
}{
{model.CapabilityText, model.APIChat, model.KindText},
{model.CapabilityImageGenerate, model.APIImages, model.KindImage},
{model.CapabilityImageEdit, model.APIImagesEdits, model.KindImage},
} {
var count int64
if err := db.Table("active_routes").Where("capability = ?", candidate.capability).Count(&count).Error; err != nil {
t.Fatal(err)
}
if count == 0 {
return candidate.capability, candidate.apiType, candidate.kind
}
}
t.Skip("router integration test requires an unbound capability in the isolated database")
return "", "", ""
}
+279
View File
@@ -0,0 +1,279 @@
// 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}
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 {
Reserve(context.Context, ReservationRequest) (Reservation, error)
Record(context.Context, Reservation, CircuitObservation) (bool, error)
}
type ReservationRequest struct {
RoutePoolMemberID uint64
Owner string
LeaseDuration time.Duration
}
type Reservation struct {
RoutePoolMemberID uint64
Token string
State CircuitState
}
type CircuitObservation struct {
Retryable bool
Succeeded bool
ErrorCode string
}
+184
View File
@@ -0,0 +1,184 @@
package router
import (
"encoding/json"
"errors"
"reflect"
"slices"
"strings"
"testing"
"git.ilapage.cn/OPC/chorus/internal/core/model"
"git.ilapage.cn/OPC/chorus/internal/core/provider"
)
type sequenceRandom struct {
values []uint64
index int
}
func (r *sequenceRandom) Uint64n(limit uint64) uint64 {
if r.index >= len(r.values) {
return 0
}
value := r.values[r.index]
r.index++
return value
}
func validSnapshot() RouteSnapshot {
return RouteSnapshot{
Capability: model.CapabilityText,
RoutePoolID: 11,
RoutePoolVersion: 3,
PromptTemplateID: 22,
PromptTemplateKey: "text-default",
PromptTemplateVersion: 7,
MaxFailover: 2,
Members: []MemberSnapshot{
{RoutePoolMemberID: 1, ProviderModelID: 101, Weight: 1, FailureThreshold: 3, OpenSeconds: 60, HalfOpenMax: 1},
{RoutePoolMemberID: 2, ProviderModelID: 102, Weight: 3, FailureThreshold: 3, OpenSeconds: 60, HalfOpenMax: 2},
{RoutePoolMemberID: 3, ProviderModelID: 103, Weight: 6, FailureThreshold: 3, OpenSeconds: 60, HalfOpenMax: 1},
},
}
}
func availableState(memberID uint64) MemberState {
return MemberState{RoutePoolMemberID: memberID, Enabled: true, ProviderEnabled: true, ModelEnabled: true, SupportsCapability: true, CircuitState: CircuitClosed}
}
func TestRouteSnapshotEncodesOnlyRoutingMetadata(t *testing.T) {
snapshot := validSnapshot()
encoded, err := snapshot.Encode()
if err != nil {
t.Fatal(err)
}
for _, forbidden := range []string{"api_key", "credential", "base_url", "nonce", "ciphertext"} {
if strings.Contains(string(encoded), forbidden) {
t.Fatalf("snapshot contains forbidden %q: %s", forbidden, encoded)
}
}
var decoded RouteSnapshot
if err := json.Unmarshal(encoded, &decoded); err != nil || !reflect.DeepEqual(decoded, snapshot) {
t.Fatalf("snapshot round trip = %#v, %v", decoded, err)
}
generation := &model.Generation{}
if err := ApplySnapshot(generation, snapshot); err != nil {
t.Fatal(err)
}
if generation.RoutePoolID == nil || *generation.RoutePoolID != snapshot.RoutePoolID || generation.RoutePoolVersion == nil || *generation.RoutePoolVersion != snapshot.RoutePoolVersion || generation.PromptTemplateID == nil || *generation.PromptTemplateID != snapshot.PromptTemplateID || !slices.Equal(generation.RouteSnapshot, encoded) {
t.Fatalf("generation route fields = %#v", generation)
}
}
func TestRouteSnapshotRejectsDuplicateModelsAndInvalidMembers(t *testing.T) {
snapshot := validSnapshot()
snapshot.Members[1].ProviderModelID = snapshot.Members[0].ProviderModelID
if err := snapshot.Validate(); !errors.Is(err, ErrInvalidSnapshot) {
t.Fatalf("duplicate model error = %v", err)
}
snapshot = validSnapshot()
snapshot.Members[0].HalfOpenMax = 0
if err := snapshot.Validate(); !errors.Is(err, ErrInvalidSnapshot) {
t.Fatalf("invalid half-open limit error = %v", err)
}
}
func TestCapabilityForSubmission(t *testing.T) {
tests := []struct {
kind model.GenerationKind
hasInputs bool
want model.Capability
}{
{model.KindText, false, model.CapabilityText},
{model.KindImage, false, model.CapabilityImageGenerate},
{model.KindImage, true, model.CapabilityImageEdit},
}
for _, test := range tests {
got, err := CapabilityForSubmission(test.kind, test.hasInputs)
if err != nil || got != test.want {
t.Fatalf("CapabilityForSubmission(%q, %t) = %q, %v; want %q", test.kind, test.hasInputs, got, err, test.want)
}
}
if _, err := CapabilityForSubmission("invalid", false); !errors.Is(err, ErrRouteNotConfigured) {
t.Fatalf("invalid capability error = %v", err)
}
}
func TestSelectCandidatesUsesWeightedSamplingWithoutReplacement(t *testing.T) {
snapshot := validSnapshot()
states := []MemberState{availableState(1), availableState(2), availableState(3)}
selected, err := SelectCandidates(snapshot, states, &sequenceRandom{values: []uint64{9, 0, 0}})
if err != nil {
t.Fatal(err)
}
got := []uint64{selected[0].RoutePoolMemberID, selected[1].RoutePoolMemberID, selected[2].RoutePoolMemberID}
if want := []uint64{3, 1, 2}; !slices.Equal(got, want) {
t.Fatalf("weighted candidates = %v, want %v", got, want)
}
snapshot.MaxFailover = 0
selected, err = SelectCandidates(snapshot, states, &sequenceRandom{values: []uint64{1}})
if err != nil || len(selected) != 1 {
t.Fatalf("max_failover=0 candidates = %#v, %v", selected, err)
}
}
func TestSelectCandidatesSkipsDisabledOpenAndCapabilityMismatchMembers(t *testing.T) {
snapshot := validSnapshot()
states := []MemberState{
availableState(1),
{RoutePoolMemberID: 2, Enabled: true, ProviderEnabled: true, ModelEnabled: true, SupportsCapability: true, CircuitState: CircuitOpen},
{RoutePoolMemberID: 3, Enabled: true, ProviderEnabled: true, ModelEnabled: true, SupportsCapability: false, CircuitState: CircuitClosed},
}
selected, err := SelectCandidates(snapshot, states, &sequenceRandom{values: []uint64{0}})
if err != nil || len(selected) != 1 || selected[0].RoutePoolMemberID != 1 {
t.Fatalf("eligible candidates = %#v, %v", selected, err)
}
if _, err := SelectCandidates(snapshot, states[1:], &sequenceRandom{values: []uint64{0}}); !errors.Is(err, ErrRouteUnavailable) {
t.Fatalf("unavailable route error = %v", err)
}
}
func TestCandidateForAttemptReturnsStableExhaustionCode(t *testing.T) {
candidates := validSnapshot().Members[:1]
member, err := CandidateForAttempt(candidates, 0)
if err != nil || member.RoutePoolMemberID != candidates[0].RoutePoolMemberID {
t.Fatalf("first candidate = %#v, %v", member, err)
}
if _, err := CandidateForAttempt(candidates, 1); !errors.Is(err, ErrFailoverExhausted) || err.Error() != string(CodeFailoverExhausted) {
t.Fatalf("exhausted candidate error = %v", err)
}
}
func TestHalfOpenEligibilityRespectsProbeLimit(t *testing.T) {
state := availableState(1)
state.CircuitState = CircuitHalfOpen
state.HalfOpenMax = 2
for inFlight, want := range map[uint16]bool{0: true, 1: true, 2: false, 3: false} {
state.HalfOpenInFlight = inFlight
if got := state.Eligible(); got != want {
t.Errorf("half-open probes %d eligible = %t, want %t", inFlight, got, want)
}
}
}
func TestFailureActionPreservesRetryableRedLines(t *testing.T) {
tests := []struct {
class provider.FailureClass
want FailureAction
}{
{provider.FailureRateLimited, FailureTryNext},
{provider.FailureServer, FailureTryNext},
{provider.FailureTimeout, FailureTryNext},
{provider.FailureConnection, FailureTryNext},
{provider.FailureBadRequest, FailureStop},
{provider.FailureUnauthorized, FailureStop},
{provider.FailurePolicyRejected, FailureStop},
}
for _, test := range tests {
if got := ActionForFailure(test.class); got != test.want {
t.Errorf("ActionForFailure(%s) = %s, want %s", test.class, got, test.want)
}
}
}