Build AI gateway control plane and admin UI

2 files changed, 141 insertions, 0 deletions

diff --git a/internal/routing/router.go b/internal/routing/router.go
new file mode 100644
index 0000000..53e5261
--- /dev/null
+++ b/internal/routing/router.go
@@ -0,0 +1,81 @@
+package routing
+
+import (
+	"errors"
+	"sort"
+	"strconv"
+	"sync"
+	"sync/atomic"
+
+	"aigw/internal/catalog"
+	"aigw/internal/domain"
+)
+
+var ErrNoRoute = errors.New("no compatible upstream route")
+
+type Router struct {
+	catalog  *catalog.Catalog
+	counters sync.Map
+}
+
+func New(catalog *catalog.Catalog) *Router {
+	return &Router{catalog: catalog}
+}
+
+func (r *Router) Plan(modelID string, protocol domain.Protocol) ([]domain.Route, error) {
+	model, err := r.catalog.Model(modelID)
+	if err != nil {
+		return nil, err
+	}
+	routes := make([]domain.Route, 0, len(model.Routes))
+	for _, route := range model.Routes {
+		if route.Provider.Protocol == protocol {
+			routes = append(routes, route)
+		}
+	}
+	if len(routes) == 0 {
+		return nil, ErrNoRoute
+	}
+
+	sort.SliceStable(routes, func(i, j int) bool { return routes[i].Priority < routes[j].Priority })
+	result := make([]domain.Route, 0, len(routes))
+	for start := 0; start < len(routes); {
+		end := start + 1
+		for end < len(routes) && routes[end].Priority == routes[start].Priority {
+			end++
+		}
+		result = append(result, r.rotate(modelID, protocol, routes[start:end])...)
+		start = end
+	}
+	return result, nil
+}
+
+func (r *Router) rotate(modelID string, protocol domain.Protocol, routes []domain.Route) []domain.Route {
+	if len(routes) < 2 {
+		return append([]domain.Route(nil), routes...)
+	}
+	key := modelID + "\x00" + string(protocol) + "\x00" + strconv.Itoa(routes[0].Priority)
+	counterValue, _ := r.counters.LoadOrStore(key, &atomic.Uint64{})
+	counter := counterValue.(*atomic.Uint64).Add(1) - 1
+
+	totalWeight := 0
+	for _, route := range routes {
+		totalWeight += route.Weight
+	}
+	position := int(counter % uint64(totalWeight))
+	selected := 0
+	for i, route := range routes {
+		if position < route.Weight {
+			selected = i
+			break
+		}
+		position -= route.Weight
+	}
+
+	result := make([]domain.Route, 0, len(routes))
+	result = append(result, routes[selected])
+	for offset := 1; offset < len(routes); offset++ {
+		result = append(result, routes[(selected+offset)%len(routes)])
+	}
+	return result
+}
diff --git a/internal/routing/router_test.go b/internal/routing/router_test.go
new file mode 100644
index 0000000..62dc656
--- /dev/null
+++ b/internal/routing/router_test.go
@@ -0,0 +1,60 @@
+package routing
+
+import (
+	"testing"
+
+	"aigw/internal/catalog"
+	"aigw/internal/config"
+	"aigw/internal/domain"
+)
+
+func TestPlanHonorsPriorityAndProtocol(t *testing.T) {
+	cfg := config.Config{
+		Providers: []config.ProviderConfig{
+			{ID: "openai-primary", Protocol: domain.ProtocolOpenAI, BaseURL: "https://one.test", APIKey: "one"},
+			{ID: "openai-fallback", Protocol: domain.ProtocolOpenAI, BaseURL: "https://two.test", APIKey: "two"},
+			{ID: "anthropic", Protocol: domain.ProtocolAnthropic, BaseURL: "https://three.test", APIKey: "three"},
+		},
+		Models: []config.ModelConfig{{
+			ID: "public/model",
+			Routes: []config.RouteConfig{
+				{Provider: "openai-fallback", UpstreamModel: "fallback", Priority: 10, Weight: 1},
+				{Provider: "anthropic", UpstreamModel: "claude", Priority: 0, Weight: 1},
+				{Provider: "openai-primary", UpstreamModel: "primary", Priority: 0, Weight: 1},
+			},
+		}},
+	}
+	router := New(catalog.New(cfg))
+	plan, err := router.Plan("public/model", domain.ProtocolOpenAI)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(plan) != 2 || plan[0].Provider.ID != "openai-primary" || plan[1].Provider.ID != "openai-fallback" {
+		t.Fatalf("unexpected plan: %+v", plan)
+	}
+}
+
+func TestPlanUsesWeightsForPrimarySelection(t *testing.T) {
+	cfg := config.Config{
+		Providers: []config.ProviderConfig{
+			{ID: "one", Protocol: domain.ProtocolOpenAI, BaseURL: "https://one.test", APIKey: "one"},
+			{ID: "two", Protocol: domain.ProtocolOpenAI, BaseURL: "https://two.test", APIKey: "two"},
+		},
+		Models: []config.ModelConfig{{ID: "public/model", Routes: []config.RouteConfig{
+			{Provider: "one", UpstreamModel: "one", Weight: 3},
+			{Provider: "two", UpstreamModel: "two", Weight: 1},
+		}}},
+	}
+	router := New(catalog.New(cfg))
+	counts := map[string]int{}
+	for range 8 {
+		plan, err := router.Plan("public/model", domain.ProtocolOpenAI)
+		if err != nil {
+			t.Fatal(err)
+		}
+		counts[plan[0].Provider.ID]++
+	}
+	if counts["one"] != 6 || counts["two"] != 2 {
+		t.Fatalf("unexpected weighted distribution: %+v", counts)
+	}
+}