feat: 实现路由核心与故障转移契约 (#21)
This commit is contained in:
@@ -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"`
|
||||
|
||||
@@ -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"}},
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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 "", "", ""
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user