mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-12 04:54:45 +00:00
feat: model registry with discovery and aggregated /v1/models
Adds modelRegistry: per-provider explicit+discovered model sets, bare-id and provider-prefixed resolution with sorted-name collision precedence, group authorization pairing (allowsModel AND allowsProvider), and a group-filtered OpenAI-compatible /v1/models catalog. Discovery is lazy and throttled per provider: newGateway runs a synchronous 5s-bounded warm fill for discovery-enabled providers (non-fatal on error), and Gateway.ServeHTTP triggers a throttled, inFlight-guarded background refresh per provider per interval (default 1h, per-provider DiscoveryInterval parsed at construction). A failed refresh keeps the previous discovered set (stale-while-error). Wires GET /v1/models into ServeHTTP (auth required, 401 envelope otherwise) and gives Gateway its adapters/registry fields, both built in newGateway. newModelRegistry's signature gained an error return beyond the task brief's literal listing: an invalid DiscoveryInterval must be a constructor error per the controller ruling, which the brief's bare *modelRegistry signature could not express.
This commit is contained in:
+29
-7
@@ -101,13 +101,13 @@ func CreateConfig() *Config {
|
||||
|
||||
// Gateway is the Traefik middleware handler.
|
||||
type Gateway struct {
|
||||
next http.Handler
|
||||
cfg *Config
|
||||
auth *authStore
|
||||
limiter *limiter
|
||||
name string
|
||||
// wired in later tasks: registry *modelRegistry, adapters
|
||||
// map[string]providerAdapter
|
||||
next http.Handler
|
||||
cfg *Config
|
||||
auth *authStore
|
||||
limiter *limiter
|
||||
registry *modelRegistry
|
||||
adapters map[string]providerAdapter
|
||||
name string
|
||||
}
|
||||
|
||||
// New creates the middleware. NOTE: no tail call — Yaegi zeroes
|
||||
@@ -143,6 +143,21 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str
|
||||
lim.logf = g.errorf
|
||||
g.limiter = lim
|
||||
|
||||
adapters, err := buildAdapters(config)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
g.adapters = adapters
|
||||
|
||||
registry, err := newModelRegistry(adapters, config, g.errorf)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
warmCtx, warmCancel := context.WithTimeout(ctx, warmFillTimeout)
|
||||
defer warmCancel()
|
||||
registry.warmFill(warmCtx)
|
||||
g.registry = registry
|
||||
|
||||
setPricingWarnFn(func(msg string) {
|
||||
if msg == warnCapMessage {
|
||||
g.logf("%s", msg)
|
||||
@@ -226,6 +241,13 @@ func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
sw := &statusTrackingWriter{ResponseWriter: w}
|
||||
defer recoverPanic(sw, g)
|
||||
g.auth.maybeReload()
|
||||
g.registry.maybeRefresh(r.Context())
|
||||
|
||||
if r.Method == http.MethodGet && r.URL.Path == "/v1/models" {
|
||||
g.handleModels(sw, r)
|
||||
return
|
||||
}
|
||||
|
||||
if g.cfg.PassthroughUnknown {
|
||||
g.next.ServeHTTP(sw, r)
|
||||
return
|
||||
|
||||
+388
@@ -0,0 +1,388 @@
|
||||
package traefikllmgateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sort"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// defaultDiscoveryInterval is the refresh interval applied to a
|
||||
// discovery-enabled provider whose ProviderConfig.DiscoveryInterval is
|
||||
// empty.
|
||||
const defaultDiscoveryInterval = time.Hour
|
||||
|
||||
// warmFillTimeout bounds newGateway's synchronous first discovery fill, per
|
||||
// provider: a slow or unreachable upstream must not stall plugin
|
||||
// construction indefinitely.
|
||||
const warmFillTimeout = 5 * time.Second
|
||||
|
||||
// backgroundRefreshTimeout bounds one background refresh goroutine spawned
|
||||
// by maybeRefresh. Not specified by the task brief; chosen generously
|
||||
// enough for a provider's models-list call while still bounding a leaked
|
||||
// goroutine's lifetime if an upstream hangs.
|
||||
const backgroundRefreshTimeout = 30 * time.Second
|
||||
|
||||
// errModelUnknown is returned by modelRegistry.resolve when id does not
|
||||
// match any provider's known model set (explicit config plus the last
|
||||
// successful discovery fetch) — 404 semantics for the caller.
|
||||
var errModelUnknown = errors.New("llmgateway: unknown model")
|
||||
|
||||
// errModelDenied is returned by modelRegistry.resolve when id is a known
|
||||
// model but the requesting group is not authorized for it — 403 semantics
|
||||
// for the caller. Distinct from errModelUnknown so the caller can tell "no
|
||||
// such model" from "that model exists, you cannot use it".
|
||||
var errModelDenied = errors.New("llmgateway: model access denied")
|
||||
|
||||
// providerState tracks one provider's known model ids and discovery
|
||||
// bookkeeping. explicit is set once at construction from
|
||||
// ProviderConfig.Models and never mutated afterward; discovered is
|
||||
// stale-while-error — a failed refresh leaves the previous discovered set
|
||||
// in place rather than clearing it. mu guards every mutable field below it.
|
||||
type providerState struct {
|
||||
lastRefresh time.Time
|
||||
explicit map[string]bool
|
||||
discovered map[string]bool
|
||||
interval time.Duration
|
||||
mu sync.Mutex
|
||||
discoveryEnabled bool
|
||||
inFlight bool
|
||||
}
|
||||
|
||||
// hasModel reports whether id is in this provider's explicit or discovered
|
||||
// model set.
|
||||
func (st *providerState) hasModel(id string) bool {
|
||||
st.mu.Lock()
|
||||
defer st.mu.Unlock()
|
||||
return st.explicit[id] || st.discovered[id]
|
||||
}
|
||||
|
||||
// knownIDs returns the sorted union of this provider's explicit and
|
||||
// discovered model ids.
|
||||
func (st *providerState) knownIDs() []string {
|
||||
st.mu.Lock()
|
||||
defer st.mu.Unlock()
|
||||
set := make(map[string]bool, len(st.explicit)+len(st.discovered))
|
||||
for id := range st.explicit {
|
||||
set[id] = true
|
||||
}
|
||||
for id := range st.discovered {
|
||||
set[id] = true
|
||||
}
|
||||
ids := make([]string, 0, len(set))
|
||||
for id := range set {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
sort.Strings(ids)
|
||||
return ids
|
||||
}
|
||||
|
||||
// tryBeginRefresh reports whether now is far enough past lastRefresh (or
|
||||
// this is the first refresh) to start a new one, and if so marks the
|
||||
// provider inFlight so a concurrent caller cannot start a second one. The
|
||||
// caller must pair a true result with a later finishRefresh call.
|
||||
func (st *providerState) tryBeginRefresh(now time.Time) bool {
|
||||
st.mu.Lock()
|
||||
defer st.mu.Unlock()
|
||||
if st.inFlight || now.Sub(st.lastRefresh) < st.interval {
|
||||
return false
|
||||
}
|
||||
st.inFlight = true
|
||||
return true
|
||||
}
|
||||
|
||||
// finishRefresh records the outcome of a refresh attempt started by
|
||||
// tryBeginRefresh (or, for the synchronous warm fill, run unconditionally).
|
||||
// lastRefresh always advances to now, on success or failure alike — that is
|
||||
// what throttles a failing provider to one attempt per interval instead of
|
||||
// retrying on every call. On success (err == nil) ids replaces the
|
||||
// discovered set; on failure the previous discovered set is kept
|
||||
// (stale-while-error).
|
||||
func (st *providerState) finishRefresh(now time.Time, ids []string, err error) {
|
||||
st.mu.Lock()
|
||||
defer st.mu.Unlock()
|
||||
st.inFlight = false
|
||||
st.lastRefresh = now
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
discovered := make(map[string]bool, len(ids))
|
||||
for _, id := range ids {
|
||||
discovered[id] = true
|
||||
}
|
||||
st.discovered = discovered
|
||||
}
|
||||
|
||||
// modelRegistry aggregates every configured provider's explicit and
|
||||
// discovered model ids into the gateway's model-resolution and /v1/models
|
||||
// listing surfaces. adapters, providerNames, and states are all fixed at
|
||||
// construction; only providerState's own fields mutate afterward (each
|
||||
// guarded by its own mutex), so modelRegistry itself needs no lock beyond
|
||||
// warnedMu for the collision-log dedup below.
|
||||
type modelRegistry struct {
|
||||
adapters map[string]providerAdapter
|
||||
states map[string]*providerState
|
||||
log func(string, ...any)
|
||||
nowFn func() time.Time
|
||||
warned map[string]bool
|
||||
providerNames []string
|
||||
warnedMu sync.Mutex
|
||||
}
|
||||
|
||||
// newModelRegistry builds a modelRegistry from adapters and cfg's matching
|
||||
// ProviderConfig entries. It does not perform discovery itself — the
|
||||
// caller (newGateway) runs the synchronous first fill separately via
|
||||
// warmFill, then ServeHTTP entry keeps it fresh via maybeRefresh. A
|
||||
// provider's ProviderConfig.DiscoveryInterval that fails time.ParseDuration
|
||||
// is a constructor error; an empty one defaults to defaultDiscoveryInterval.
|
||||
func newModelRegistry(adapters map[string]providerAdapter, cfg *Config, log func(string, ...any)) (*modelRegistry, error) {
|
||||
m := &modelRegistry{
|
||||
adapters: adapters,
|
||||
states: make(map[string]*providerState, len(adapters)),
|
||||
log: log,
|
||||
nowFn: time.Now,
|
||||
warned: make(map[string]bool),
|
||||
providerNames: make([]string, 0, len(adapters)),
|
||||
}
|
||||
for name := range adapters {
|
||||
m.providerNames = append(m.providerNames, name)
|
||||
}
|
||||
sort.Strings(m.providerNames)
|
||||
|
||||
for _, name := range m.providerNames {
|
||||
pc := cfg.Providers[name]
|
||||
if pc == nil {
|
||||
pc = &ProviderConfig{}
|
||||
}
|
||||
|
||||
interval := defaultDiscoveryInterval
|
||||
if pc.DiscoveryInterval != "" {
|
||||
d, err := time.ParseDuration(pc.DiscoveryInterval)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("llmgateway: provider %q: invalid discoveryInterval %q: %w", name, pc.DiscoveryInterval, err)
|
||||
}
|
||||
interval = d
|
||||
}
|
||||
|
||||
explicit := make(map[string]bool, len(pc.Models))
|
||||
for _, id := range pc.Models {
|
||||
explicit[id] = true
|
||||
}
|
||||
|
||||
m.states[name] = &providerState{
|
||||
explicit: explicit,
|
||||
discovered: make(map[string]bool),
|
||||
discoveryEnabled: pc.Discovery,
|
||||
interval: interval,
|
||||
}
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// now returns the registry's current time, via nowFn — overridable in
|
||||
// tests to exercise refresh throttling deterministically.
|
||||
func (m *modelRegistry) now() time.Time {
|
||||
return m.nowFn()
|
||||
}
|
||||
|
||||
// warmFill performs newGateway's synchronous first discovery fill: for
|
||||
// every discovery-enabled provider, it fetches listModels once, bounded by
|
||||
// warmFillTimeout, and records the result via finishRefresh. A fetch error
|
||||
// is logged and otherwise non-fatal — construction still succeeds, and the
|
||||
// provider's known model set is whatever its explicit config already
|
||||
// provides (empty, if it configured neither explicit models nor a
|
||||
// reachable discovery endpoint) until the next maybeRefresh window.
|
||||
func (m *modelRegistry) warmFill(ctx context.Context) {
|
||||
for _, name := range m.providerNames {
|
||||
st := m.states[name]
|
||||
if !st.discoveryEnabled {
|
||||
continue
|
||||
}
|
||||
fctx, cancel := context.WithTimeout(ctx, warmFillTimeout)
|
||||
ids, err := m.adapters[name].listModels(fctx)
|
||||
cancel()
|
||||
st.finishRefresh(m.now(), ids, err)
|
||||
if err != nil {
|
||||
m.log("model registry: initial discovery for provider %q failed: %v", name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// maybeRefresh is called at every ServeHTTP entry. For each discovery-
|
||||
// enabled provider whose last refresh is older than its interval (and that
|
||||
// is not already mid-refresh), it spawns exactly one background goroutine
|
||||
// to pull listModels and merge the result. The spawned goroutine
|
||||
// deliberately does not inherit ctx's cancellation: ctx is the incoming
|
||||
// request's context, which net/http cancels once ServeHTTP returns — using
|
||||
// it here would abort the background refresh moments after it starts on
|
||||
// every request, so the goroutine gets its own backgroundRefreshTimeout
|
||||
// budget instead. ctx is still consulted before spawning: a request whose
|
||||
// context is already done skips starting new background work for it.
|
||||
func (m *modelRegistry) maybeRefresh(ctx context.Context) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return
|
||||
}
|
||||
now := m.now()
|
||||
for _, name := range m.providerNames {
|
||||
st := m.states[name]
|
||||
if !st.discoveryEnabled || !st.tryBeginRefresh(now) {
|
||||
continue
|
||||
}
|
||||
adapter := m.adapters[name]
|
||||
go m.refreshProvider(name, st, adapter) //nolint:gosec // G118: refreshProvider intentionally uses its own context.Background()-derived timeout, not r.Context(), per this function's doc comment above — the request context is canceled once ServeHTTP returns and would abort the background fetch moments after it starts
|
||||
}
|
||||
}
|
||||
|
||||
// refreshProvider runs one background discovery fetch for a provider
|
||||
// already marked inFlight by tryBeginRefresh, and records its outcome via
|
||||
// finishRefresh. Split out from maybeRefresh so the goroutine closes over
|
||||
// named parameters, not loop variables.
|
||||
func (m *modelRegistry) refreshProvider(name string, st *providerState, adapter providerAdapter) {
|
||||
fctx, cancel := context.WithTimeout(context.Background(), backgroundRefreshTimeout)
|
||||
defer cancel()
|
||||
ids, err := adapter.listModels(fctx)
|
||||
st.finishRefresh(m.now(), ids, err)
|
||||
if err != nil {
|
||||
m.log("model registry: discovery refresh for provider %q failed: %v", name, err)
|
||||
}
|
||||
}
|
||||
|
||||
// splitConfiguredProvider reports whether id has "prefix/rest" form where
|
||||
// prefix names a currently configured provider. Per ruling (a): a slash in
|
||||
// id is only ever a provider-prefix separator when the part before it is a
|
||||
// real, configured provider name; otherwise (including when id has no
|
||||
// slash, or its prefix names no configured provider) the whole string is a
|
||||
// bare model id — this is what lets an upstream model id that itself
|
||||
// contains a slash (e.g. "uni/deepseek-v4-flash-0731" on a provider that
|
||||
// happens not to be named "uni") resolve correctly as a bare id.
|
||||
func (m *modelRegistry) splitConfiguredProvider(id string) (providerName, rest string, ok bool) {
|
||||
for i := 0; i < len(id); i++ {
|
||||
if id[i] != '/' {
|
||||
continue
|
||||
}
|
||||
prefix, tail := id[:i], id[i+1:]
|
||||
if _, known := m.adapters[prefix]; known {
|
||||
return prefix, tail, true
|
||||
}
|
||||
return "", "", false
|
||||
}
|
||||
return "", "", false
|
||||
}
|
||||
|
||||
// bareWinner returns the provider that owns bare id when more than one
|
||||
// provider's known model set contains it: the first, in sorted
|
||||
// provider-name order (ruling (g)).
|
||||
func (m *modelRegistry) bareWinner(id string) (string, bool) {
|
||||
for _, name := range m.providerNames {
|
||||
if m.states[name].hasModel(id) {
|
||||
return name, true
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
// resolve maps a client-requested model id to the adapter that serves it.
|
||||
// It returns the adapter, the upstream model id to send that adapter (the
|
||||
// id form the provider itself knows, with any gateway-side provider prefix
|
||||
// stripped), and the canonical "provider/upstreamModel" id. The original,
|
||||
// exactly-as-requested id the caller passed in is not returned separately —
|
||||
// the caller already holds it, in the id argument itself.
|
||||
//
|
||||
// id in "provider/model" form (ruling (a): only when "provider" names a
|
||||
// configured provider) resolves directly against that provider. Any other
|
||||
// id is a bare id, resolved against the first configured provider (sorted
|
||||
// name order) whose known model set contains it — ruling (g)'s collision
|
||||
// rule. Either way, authorization requires both grp.allowsModel and
|
||||
// grp.allowsProvider for the resolved provider (ruling (c)); a model that
|
||||
// exists but fails authorization returns errModelDenied, distinct from
|
||||
// errModelUnknown for a model no configured provider knows at all.
|
||||
func (m *modelRegistry) resolve(id string, grp *group) (providerAdapter, string, string, error) {
|
||||
if providerName, rest, ok := m.splitConfiguredProvider(id); ok {
|
||||
return m.resolveAgainst(providerName, rest, id, grp)
|
||||
}
|
||||
providerName, ok := m.bareWinner(id)
|
||||
if !ok {
|
||||
return nil, "", "", errModelUnknown
|
||||
}
|
||||
return m.resolveAgainst(providerName, id, id, grp)
|
||||
}
|
||||
|
||||
// resolveAgainst finishes resolve for a providerName already chosen (either
|
||||
// the explicit prefix or the bare-id collision winner): it checks
|
||||
// upstreamModel is actually known to that provider, then checks
|
||||
// authorization, in that order — so a real model behind a provider the
|
||||
// group cannot use reports errModelDenied, not errModelUnknown.
|
||||
// requestedID is the id form passed to grp.allowsModel; it is the
|
||||
// provider-prefixed form for a direct request, or the bare id for a
|
||||
// bare-id request — allowsModel already checks both a pattern's full-string
|
||||
// and bare-suffix match, so either form is checked correctly either way.
|
||||
func (m *modelRegistry) resolveAgainst(providerName, upstreamModel, requestedID string, grp *group) (providerAdapter, string, string, error) {
|
||||
if !m.states[providerName].hasModel(upstreamModel) {
|
||||
return nil, "", "", errModelUnknown
|
||||
}
|
||||
if !grp.allowsModel(requestedID) || !grp.allowsProvider(providerName) {
|
||||
return nil, "", "", errModelDenied
|
||||
}
|
||||
return m.adapters[providerName], upstreamModel, providerName + "/" + upstreamModel, nil
|
||||
}
|
||||
|
||||
// warnCollisionOnce logs, at most once per colliding bare id for this
|
||||
// registry's lifetime, that provs (sorted) all provide id and winner owns
|
||||
// its bare form.
|
||||
func (m *modelRegistry) warnCollisionOnce(id, winner string, provs []string) {
|
||||
m.warnedMu.Lock()
|
||||
defer m.warnedMu.Unlock()
|
||||
if m.warned[id] {
|
||||
return
|
||||
}
|
||||
m.warned[id] = true
|
||||
m.log("model registry: model id %q is provided by multiple providers %v; %q wins the bare id", id, provs, winner)
|
||||
}
|
||||
|
||||
// listFor returns grp's visible model catalog as OpenAI-compatible model
|
||||
// objects ({"id","object":"model","owned_by"}), sorted by id. A bare id
|
||||
// owned by only one provider is listed once, bare. A bare id owned by more
|
||||
// than one provider (ruling (g)/(h)) is listed once bare — under its
|
||||
// collision winner — plus once more per owning provider in "provider/id"
|
||||
// form, so every provider's copy stays reachable through explicit
|
||||
// addressing even when it lost the bare-id collision. Each listed entry is
|
||||
// independently filtered by grp.allowsModel and grp.allowsProvider for the
|
||||
// provider it names.
|
||||
func (m *modelRegistry) listFor(grp *group) []map[string]any {
|
||||
owners := make(map[string][]string)
|
||||
for _, name := range m.providerNames {
|
||||
for _, id := range m.states[name].knownIDs() {
|
||||
owners[id] = append(owners[id], name) // m.providerNames is sorted, so owners[id] accumulates in sorted order
|
||||
}
|
||||
}
|
||||
|
||||
out := make([]map[string]any, 0, len(owners))
|
||||
for id, provs := range owners {
|
||||
winner := provs[0]
|
||||
if grp.allowsModel(id) && grp.allowsProvider(winner) {
|
||||
out = append(out, modelObject(id, winner))
|
||||
}
|
||||
if len(provs) < 2 {
|
||||
continue
|
||||
}
|
||||
m.warnCollisionOnce(id, winner, provs)
|
||||
for _, p := range provs {
|
||||
pid := p + "/" + id
|
||||
if grp.allowsModel(pid) && grp.allowsProvider(p) {
|
||||
out = append(out, modelObject(pid, p))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
sort.Slice(out, func(i, j int) bool {
|
||||
return out[i]["id"].(string) < out[j]["id"].(string) //nolint:forcetypeassert // modelObject always sets id to a string
|
||||
})
|
||||
return out
|
||||
}
|
||||
|
||||
// modelObject builds one OpenAI-compatible model list entry.
|
||||
func modelObject(id, ownedBy string) map[string]any {
|
||||
return map[string]any{"id": id, "object": "model", "owned_by": ownedBy}
|
||||
}
|
||||
@@ -0,0 +1,731 @@
|
||||
package traefikllmgateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// fakeAdapter is a providerAdapter test double whose listModels behavior is
|
||||
// scripted per test via listModelsFn. Its other methods are never exercised
|
||||
// by registry.go, so they return harmless zero values.
|
||||
type fakeAdapter struct {
|
||||
listModelsFn func(context.Context) ([]string, error)
|
||||
adapterName string
|
||||
}
|
||||
|
||||
func newFakeAdapter(name string) *fakeAdapter {
|
||||
return &fakeAdapter{adapterName: name}
|
||||
}
|
||||
|
||||
func (f *fakeAdapter) name() string { return f.adapterName }
|
||||
func (f *fakeAdapter) typeName() string { return "fake" }
|
||||
func (f *fakeAdapter) base() string { return "" }
|
||||
func (f *fakeAdapter) injectAuth(*http.Request) {}
|
||||
func (f *fakeAdapter) chatCompletion(context.Context, http.ResponseWriter, map[string]any) (usage, error) {
|
||||
return usage{}, nil
|
||||
}
|
||||
func (f *fakeAdapter) embeddings(context.Context, http.ResponseWriter, map[string]any) (usage, error) {
|
||||
return usage{}, nil
|
||||
}
|
||||
func (f *fakeAdapter) listModels(ctx context.Context) ([]string, error) {
|
||||
if f.listModelsFn == nil {
|
||||
return nil, nil
|
||||
}
|
||||
return f.listModelsFn(ctx)
|
||||
}
|
||||
|
||||
// recordingLog is a thread-safe log func for tests, capturing formatted
|
||||
// lines instead of writing to stderr.
|
||||
type recordingLog struct {
|
||||
lines []string
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func (r *recordingLog) fn(format string, args ...any) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.lines = append(r.lines, fmt.Sprintf(format, args...))
|
||||
}
|
||||
|
||||
func (r *recordingLog) count() int {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return len(r.lines)
|
||||
}
|
||||
|
||||
// waitUntil polls cond every millisecond until it returns true, failing the
|
||||
// test if timeout elapses first. Used to synchronize against background
|
||||
// refresh goroutines without a fixed sleep.
|
||||
func waitUntil(t *testing.T, timeout time.Duration, cond func() bool) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(timeout)
|
||||
for {
|
||||
if cond() {
|
||||
return
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatal("condition not met before timeout")
|
||||
}
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
// allowAllGroup has no provider/model restrictions.
|
||||
func allowAllGroup() *group { return &group{name: "all"} }
|
||||
|
||||
// --- resolve: prefix vs bare-id disambiguation (ruling a) ---
|
||||
|
||||
func TestModelRegistry_Resolve_ProviderPrefixed_ResolvesDirect(t *testing.T) {
|
||||
t.Parallel()
|
||||
adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{"openai": {Models: []string{"gpt-test"}}}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
adapter, upstreamModel, canonical, err := reg.resolve("openai/gpt-test", allowAllGroup())
|
||||
if err != nil {
|
||||
t.Fatalf("resolve: %v", err)
|
||||
}
|
||||
if adapter != adapters["openai"] {
|
||||
t.Errorf("adapter = %v, want openai adapter", adapter)
|
||||
}
|
||||
if upstreamModel != "gpt-test" {
|
||||
t.Errorf("upstreamModel = %q, want %q", upstreamModel, "gpt-test")
|
||||
}
|
||||
if canonical != "openai/gpt-test" {
|
||||
t.Errorf("canonical = %q, want %q", canonical, "openai/gpt-test")
|
||||
}
|
||||
}
|
||||
|
||||
// TestModelRegistry_Resolve_BareIDContainingSlash_UnconfiguredPrefix_ResolvesAsBareID
|
||||
// is the real-case regression from ruling (a): an upstream model id
|
||||
// ("uni/deepseek-v4-flash-0731") that itself contains a slash must resolve
|
||||
// as a bare id against whichever configured provider lists it, because
|
||||
// "uni" itself is not a configured provider name.
|
||||
func TestModelRegistry_Resolve_BareIDContainingSlash_UnconfiguredPrefix_ResolvesAsBareID(t *testing.T) {
|
||||
t.Parallel()
|
||||
adapters := map[string]providerAdapter{
|
||||
"openai": newFakeAdapter("openai"),
|
||||
"acme": newFakeAdapter("acme"),
|
||||
}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{
|
||||
"openai": {Models: []string{"gpt-test"}},
|
||||
"acme": {Models: []string{"uni/deepseek-v4-flash-0731"}},
|
||||
}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
const id = "uni/deepseek-v4-flash-0731"
|
||||
adapter, upstreamModel, canonical, err := reg.resolve(id, allowAllGroup())
|
||||
if err != nil {
|
||||
t.Fatalf("resolve: %v", err)
|
||||
}
|
||||
if adapter != adapters["acme"] {
|
||||
t.Errorf("adapter = %v, want acme adapter", adapter)
|
||||
}
|
||||
if upstreamModel != id {
|
||||
t.Errorf("upstreamModel = %q, want unchanged %q", upstreamModel, id)
|
||||
}
|
||||
if canonical != "acme/"+id {
|
||||
t.Errorf("canonical = %q, want %q", canonical, "acme/"+id)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelRegistry_Resolve_UnknownModel_ReturnsErrModelUnknown(t *testing.T) {
|
||||
t.Parallel()
|
||||
adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{"openai": {Models: []string{"gpt-test"}}}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
if _, _, _, err := reg.resolve("does-not-exist", allowAllGroup()); err != errModelUnknown {
|
||||
t.Errorf("err = %v, want errModelUnknown", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelRegistry_Resolve_EmptyProviderModelSet_ResolvesNothing(t *testing.T) {
|
||||
t.Parallel()
|
||||
adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{"openai": {}}} // no explicit models, discovery off
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
if _, _, _, err := reg.resolve("openai/anything", allowAllGroup()); err != errModelUnknown {
|
||||
t.Errorf("err = %v, want errModelUnknown", err)
|
||||
}
|
||||
if _, _, _, err := reg.resolve("anything", allowAllGroup()); err != errModelUnknown {
|
||||
t.Errorf("err = %v, want errModelUnknown", err)
|
||||
}
|
||||
}
|
||||
|
||||
// --- resolve: authorization pairing (ruling c) ---
|
||||
|
||||
func TestModelRegistry_Resolve_KnownModel_ProviderDenied_ReturnsErrModelDenied(t *testing.T) {
|
||||
t.Parallel()
|
||||
adapters := map[string]providerAdapter{
|
||||
"openai": newFakeAdapter("openai"),
|
||||
"anthropic": newFakeAdapter("anthropic"),
|
||||
}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{
|
||||
"openai": {Models: []string{"gpt-test"}},
|
||||
"anthropic": {Models: []string{"claude-x"}},
|
||||
}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
// providers:[openai] restricts the provider; models:[claude-*] would
|
||||
// match claude-x on its own — the provider gate must still deny it.
|
||||
grp := &group{name: "openai-only", providers: []string{"openai"}, models: []string{"claude-*"}}
|
||||
if _, _, _, err := reg.resolve("anthropic/claude-x", grp); err != errModelDenied {
|
||||
t.Errorf("err = %v, want errModelDenied", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelRegistry_Resolve_KnownModel_ModelGlobDenied_ReturnsErrModelDenied(t *testing.T) {
|
||||
t.Parallel()
|
||||
adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{"openai": {Models: []string{"gpt-test"}}}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
grp := &group{name: "narrow", providers: []string{"openai"}, models: []string{"gpt-4*"}}
|
||||
if _, _, _, err := reg.resolve("openai/gpt-test", grp); err != errModelDenied {
|
||||
t.Errorf("err = %v, want errModelDenied", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelRegistry_Resolve_AllowedByBothProviderAndModel_Succeeds(t *testing.T) {
|
||||
t.Parallel()
|
||||
adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{"openai": {Models: []string{"gpt-test"}}}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
grp := &group{name: "ok", providers: []string{"openai"}, models: []string{"gpt-*"}}
|
||||
if _, _, _, err := reg.resolve("openai/gpt-test", grp); err != nil {
|
||||
t.Errorf("resolve: %v, want success", err)
|
||||
}
|
||||
}
|
||||
|
||||
// --- resolve: collision precedence (ruling g) ---
|
||||
|
||||
func TestModelRegistry_Resolve_BareIDCollision_PicksSortedFirstProvider(t *testing.T) {
|
||||
t.Parallel()
|
||||
adapters := map[string]providerAdapter{
|
||||
"beta": newFakeAdapter("beta"),
|
||||
"alpha": newFakeAdapter("alpha"),
|
||||
}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{
|
||||
"beta": {Models: []string{"shared"}},
|
||||
"alpha": {Models: []string{"shared"}},
|
||||
}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
adapter, _, canonical, err := reg.resolve("shared", allowAllGroup())
|
||||
if err != nil {
|
||||
t.Fatalf("resolve: %v", err)
|
||||
}
|
||||
if adapter != adapters["alpha"] {
|
||||
t.Errorf("adapter = %v, want alpha (sorted-first)", adapter)
|
||||
}
|
||||
if canonical != "alpha/shared" {
|
||||
t.Errorf("canonical = %q, want %q", canonical, "alpha/shared")
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelRegistry_Resolve_ProviderPrefixed_AlwaysReachesLosingProvider(t *testing.T) {
|
||||
t.Parallel()
|
||||
adapters := map[string]providerAdapter{
|
||||
"beta": newFakeAdapter("beta"),
|
||||
"alpha": newFakeAdapter("alpha"),
|
||||
}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{
|
||||
"beta": {Models: []string{"shared"}},
|
||||
"alpha": {Models: []string{"shared"}},
|
||||
}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
adapter, _, canonical, err := reg.resolve("beta/shared", allowAllGroup())
|
||||
if err != nil {
|
||||
t.Fatalf("resolve: %v", err)
|
||||
}
|
||||
if adapter != adapters["beta"] {
|
||||
t.Errorf("adapter = %v, want beta (explicit form bypasses collision winner)", adapter)
|
||||
}
|
||||
if canonical != "beta/shared" {
|
||||
t.Errorf("canonical = %q, want %q", canonical, "beta/shared")
|
||||
}
|
||||
}
|
||||
|
||||
// --- newModelRegistry construction ---
|
||||
|
||||
func TestNewModelRegistry_InvalidDiscoveryInterval_ReturnsConstructorError(t *testing.T) {
|
||||
t.Parallel()
|
||||
adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{"openai": {DiscoveryInterval: "not-a-duration"}}}
|
||||
if _, err := newModelRegistry(adapters, cfg, func(string, ...any) {}); err == nil {
|
||||
t.Fatal("want constructor error for an invalid discoveryInterval")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewModelRegistry_EmptyDiscoveryInterval_DefaultsOneHour(t *testing.T) {
|
||||
t.Parallel()
|
||||
adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{"openai": {}}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
if got := reg.states["openai"].interval; got != defaultDiscoveryInterval {
|
||||
t.Errorf("interval = %v, want %v", got, defaultDiscoveryInterval)
|
||||
}
|
||||
}
|
||||
|
||||
// --- warmFill: synchronous initial discovery (ruling e) ---
|
||||
|
||||
func TestModelRegistry_WarmFill_MergesExplicitAndDiscovered(t *testing.T) {
|
||||
t.Parallel()
|
||||
fa := newFakeAdapter("openai")
|
||||
fa.listModelsFn = func(context.Context) ([]string, error) { return []string{"gpt-b", "gpt-a"}, nil }
|
||||
adapters := map[string]providerAdapter{"openai": fa}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{"openai": {Models: []string{"gpt-a"}, Discovery: true}}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
reg.warmFill(context.Background())
|
||||
|
||||
got := reg.states["openai"].knownIDs()
|
||||
want := []string{"gpt-a", "gpt-b"}
|
||||
if len(got) != len(want) || got[0] != want[0] || got[1] != want[1] {
|
||||
t.Errorf("knownIDs = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelRegistry_WarmFill_DiscoveryError_NonFatal_LogsAndLeavesExplicitModels(t *testing.T) {
|
||||
t.Parallel()
|
||||
fa := newFakeAdapter("openai")
|
||||
fa.listModelsFn = func(context.Context) ([]string, error) { return nil, fmt.Errorf("upstream down") }
|
||||
adapters := map[string]providerAdapter{"openai": fa}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{"openai": {Models: []string{"gpt-a"}, Discovery: true}}}
|
||||
rl := &recordingLog{}
|
||||
reg, err := newModelRegistry(adapters, cfg, rl.fn)
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
reg.warmFill(context.Background())
|
||||
|
||||
if !reg.states["openai"].hasModel("gpt-a") {
|
||||
t.Error("want explicit model gpt-a still known after a failed initial discovery fetch")
|
||||
}
|
||||
if rl.count() == 0 {
|
||||
t.Error("want the discovery failure logged")
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelRegistry_WarmFill_NonDiscoveryProvider_NeverCallsListModels(t *testing.T) {
|
||||
t.Parallel()
|
||||
var calls int32
|
||||
fa := newFakeAdapter("openai")
|
||||
fa.listModelsFn = func(context.Context) ([]string, error) {
|
||||
atomic.AddInt32(&calls, 1)
|
||||
return nil, nil
|
||||
}
|
||||
adapters := map[string]providerAdapter{"openai": fa}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{"openai": {Models: []string{"gpt-a"}}}} // Discovery left false
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
reg.warmFill(context.Background())
|
||||
|
||||
if got := atomic.LoadInt32(&calls); got != 0 {
|
||||
t.Errorf("listModels called %d times, want 0 for a non-discovery provider", got)
|
||||
}
|
||||
}
|
||||
|
||||
// --- maybeRefresh: throttling, in-flight guard, stale-while-error ---
|
||||
|
||||
func TestModelRegistry_MaybeRefresh_Throttled_OnePerInterval(t *testing.T) {
|
||||
t.Parallel()
|
||||
var calls int32
|
||||
fa := newFakeAdapter("openai")
|
||||
fa.listModelsFn = func(context.Context) ([]string, error) {
|
||||
atomic.AddInt32(&calls, 1)
|
||||
return []string{"m1"}, nil
|
||||
}
|
||||
adapters := map[string]providerAdapter{"openai": fa}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{"openai": {Discovery: true, DiscoveryInterval: "1h"}}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
now := time.Now()
|
||||
reg.nowFn = func() time.Time { return now }
|
||||
|
||||
reg.maybeRefresh(context.Background())
|
||||
waitUntil(t, time.Second, func() bool { return atomic.LoadInt32(&calls) == 1 })
|
||||
|
||||
// Still within the same interval: a second entry must not refetch.
|
||||
reg.maybeRefresh(context.Background())
|
||||
time.Sleep(10 * time.Millisecond) // give a wrongly-spawned goroutine a chance to run
|
||||
if got := atomic.LoadInt32(&calls); got != 1 {
|
||||
t.Errorf("listModels called %d times within one interval, want 1", got)
|
||||
}
|
||||
|
||||
now = now.Add(2 * time.Hour) // past the interval: next entry should refetch
|
||||
reg.maybeRefresh(context.Background())
|
||||
waitUntil(t, time.Second, func() bool { return atomic.LoadInt32(&calls) == 2 })
|
||||
}
|
||||
|
||||
// TestModelRegistry_MaybeRefresh_InFlightGuard_BlocksConcurrentFetch proves
|
||||
// the guard is the inFlight flag, not just the interval check: the first
|
||||
// fetch is held open (lastRefresh never advances while it's in flight), so
|
||||
// the interval check alone would spuriously allow a second concurrent
|
||||
// fetch. Only the inFlight flag prevents it.
|
||||
func TestModelRegistry_MaybeRefresh_InFlightGuard_BlocksConcurrentFetch(t *testing.T) {
|
||||
t.Parallel()
|
||||
var calls int32
|
||||
started := make(chan struct{}, 2)
|
||||
release := make(chan struct{})
|
||||
fa := newFakeAdapter("openai")
|
||||
fa.listModelsFn = func(context.Context) ([]string, error) {
|
||||
atomic.AddInt32(&calls, 1)
|
||||
started <- struct{}{}
|
||||
<-release
|
||||
return []string{"m1"}, nil
|
||||
}
|
||||
adapters := map[string]providerAdapter{"openai": fa}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{"openai": {Discovery: true, DiscoveryInterval: "1h"}}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
now := time.Now()
|
||||
reg.nowFn = func() time.Time { return now }
|
||||
|
||||
reg.maybeRefresh(context.Background()) // spawns the in-flight fetch, synchronously marks inFlight
|
||||
reg.maybeRefresh(context.Background()) // must be a no-op: inFlight already true
|
||||
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("first refresh never started")
|
||||
}
|
||||
select {
|
||||
case <-started:
|
||||
t.Fatal("a second concurrent fetch started while the first was still in flight")
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
}
|
||||
|
||||
close(release)
|
||||
waitUntil(t, time.Second, func() bool { return atomic.LoadInt32(&calls) == 1 && reg.states["openai"].hasModel("m1") })
|
||||
}
|
||||
|
||||
func TestModelRegistry_MaybeRefresh_DiscoveryError_KeepsPreviousDiscoveredSet(t *testing.T) {
|
||||
t.Parallel()
|
||||
var fail atomic.Bool
|
||||
fa := newFakeAdapter("openai")
|
||||
fa.listModelsFn = func(context.Context) ([]string, error) {
|
||||
if fail.Load() {
|
||||
return nil, fmt.Errorf("upstream unreachable")
|
||||
}
|
||||
return []string{"m1"}, nil
|
||||
}
|
||||
adapters := map[string]providerAdapter{"openai": fa}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{"openai": {Discovery: true, DiscoveryInterval: "1h"}}}
|
||||
rl := &recordingLog{}
|
||||
reg, err := newModelRegistry(adapters, cfg, rl.fn)
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
reg.warmFill(context.Background()) // succeeds, populates discovered={m1}
|
||||
if !reg.states["openai"].hasModel("m1") {
|
||||
t.Fatal("warmFill did not populate m1")
|
||||
}
|
||||
|
||||
now := time.Now().Add(2 * time.Hour)
|
||||
reg.nowFn = func() time.Time { return now }
|
||||
fail.Store(true)
|
||||
before := rl.count()
|
||||
|
||||
reg.maybeRefresh(context.Background())
|
||||
waitUntil(t, time.Second, func() bool { return rl.count() > before })
|
||||
|
||||
if !reg.states["openai"].hasModel("m1") {
|
||||
t.Error("want the previous discovered set kept after a failed refresh (stale-while-error)")
|
||||
}
|
||||
}
|
||||
|
||||
// --- listFor: group filtering, sorting, collision presentation (ruling h) ---
|
||||
|
||||
func TestModelRegistry_ListFor_GroupFiltered(t *testing.T) {
|
||||
t.Parallel()
|
||||
adapters := map[string]providerAdapter{
|
||||
"openai": newFakeAdapter("openai"),
|
||||
"anthropic": newFakeAdapter("anthropic"),
|
||||
}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{
|
||||
"openai": {Models: []string{"gpt-test"}},
|
||||
"anthropic": {Models: []string{"claude-x"}},
|
||||
}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
grp := &group{name: "openai-only", providers: []string{"openai"}}
|
||||
got := reg.listFor(grp)
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("listFor = %v, want exactly 1 entry", got)
|
||||
}
|
||||
if got[0]["id"] != "gpt-test" || got[0]["owned_by"] != "openai" || got[0]["object"] != "model" {
|
||||
t.Errorf("entry = %v, want {id:gpt-test, owned_by:openai, object:model}", got[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelRegistry_ListFor_SortedByID(t *testing.T) {
|
||||
t.Parallel()
|
||||
adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{"openai": {Models: []string{"zeta", "alpha", "mu"}}}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
got := reg.listFor(allowAllGroup())
|
||||
if len(got) != 3 {
|
||||
t.Fatalf("listFor = %v, want 3 entries", got)
|
||||
}
|
||||
ids := []string{got[0]["id"].(string), got[1]["id"].(string), got[2]["id"].(string)}
|
||||
want := []string{"alpha", "mu", "zeta"}
|
||||
for i := range want {
|
||||
if ids[i] != want[i] {
|
||||
t.Errorf("ids = %v, want %v", ids, want)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelRegistry_ListFor_Collision_ListsBareWinnerPlusAllPrefixedForms(t *testing.T) {
|
||||
t.Parallel()
|
||||
adapters := map[string]providerAdapter{
|
||||
"beta": newFakeAdapter("beta"),
|
||||
"alpha": newFakeAdapter("alpha"),
|
||||
}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{
|
||||
"beta": {Models: []string{"shared"}},
|
||||
"alpha": {Models: []string{"shared"}},
|
||||
}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
got := reg.listFor(allowAllGroup())
|
||||
byID := make(map[string]string, len(got))
|
||||
for _, e := range got {
|
||||
byID[e["id"].(string)] = e["owned_by"].(string)
|
||||
}
|
||||
want := map[string]string{"shared": "alpha", "alpha/shared": "alpha", "beta/shared": "beta"}
|
||||
if len(byID) != len(want) {
|
||||
t.Fatalf("listFor = %v, want exactly %v", byID, want)
|
||||
}
|
||||
for id, ownedBy := range want {
|
||||
if byID[id] != ownedBy {
|
||||
t.Errorf("entry %q owned_by = %q, want %q", id, byID[id], ownedBy)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestModelRegistry_ListFor_Collision_LogsWarningOnceRegardlessOfCallCount
|
||||
// covers ruling (g)'s "log warn once per colliding id": repeated listFor
|
||||
// calls over the same collision must not re-emit the warning.
|
||||
func TestModelRegistry_ListFor_Collision_LogsWarningOnceRegardlessOfCallCount(t *testing.T) {
|
||||
t.Parallel()
|
||||
adapters := map[string]providerAdapter{
|
||||
"beta": newFakeAdapter("beta"),
|
||||
"alpha": newFakeAdapter("alpha"),
|
||||
}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{
|
||||
"beta": {Models: []string{"shared"}},
|
||||
"alpha": {Models: []string{"shared"}},
|
||||
}}
|
||||
rl := &recordingLog{}
|
||||
reg, err := newModelRegistry(adapters, cfg, rl.fn)
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
reg.listFor(allowAllGroup())
|
||||
reg.listFor(allowAllGroup())
|
||||
reg.listFor(allowAllGroup())
|
||||
|
||||
if got := rl.count(); got != 1 {
|
||||
t.Errorf("collision log emitted %d times across 3 listFor calls, want exactly 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelRegistry_ListFor_Collision_LosingProviderStillReachableViaPrefixedForm(t *testing.T) {
|
||||
t.Parallel()
|
||||
adapters := map[string]providerAdapter{
|
||||
"beta": newFakeAdapter("beta"),
|
||||
"alpha": newFakeAdapter("alpha"),
|
||||
}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{
|
||||
"beta": {Models: []string{"shared"}},
|
||||
"alpha": {Models: []string{"shared"}},
|
||||
}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
grp := &group{name: "beta-only", providers: []string{"beta"}}
|
||||
got := reg.listFor(grp)
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("listFor = %v, want exactly 1 entry", got)
|
||||
}
|
||||
if got[0]["id"] != "beta/shared" || got[0]["owned_by"] != "beta" {
|
||||
t.Errorf("entry = %v, want {id:beta/shared, owned_by:beta}", got[0])
|
||||
}
|
||||
}
|
||||
|
||||
// --- end-to-end: GET /v1/models through Gateway.ServeHTTP ---
|
||||
|
||||
func TestGateway_ServeHTTP_GetV1Models_ReturnsGroupFilteredList(t *testing.T) {
|
||||
t.Parallel()
|
||||
next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})
|
||||
cfg := CreateConfig()
|
||||
cfg.Providers = map[string]*ProviderConfig{
|
||||
"openai": {Type: "openai", APIKey: "k", Models: []string{"gpt-test"}},
|
||||
"anthropic": {Type: "anthropic", APIKey: "k", Models: []string{"claude-x"}},
|
||||
}
|
||||
cfg.Groups = map[string]*GroupConfig{"openai-only": {Providers: []string{"openai"}}}
|
||||
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "openai-only", APIKey: "sk-alice"}}}
|
||||
|
||||
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
req.Header.Set("Authorization", "Bearer sk-alice")
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d, body: %s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
if ct := rec.Header().Get("Content-Type"); ct != "application/json" {
|
||||
t.Errorf("Content-Type = %q, want application/json", ct)
|
||||
}
|
||||
|
||||
var body struct {
|
||||
Object string `json:"object"`
|
||||
Data []map[string]any `json:"data"`
|
||||
}
|
||||
if err := json.NewDecoder(rec.Body).Decode(&body); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
if body.Object != "list" {
|
||||
t.Errorf("object = %q, want %q", body.Object, "list")
|
||||
}
|
||||
if len(body.Data) != 1 || body.Data[0]["id"] != "gpt-test" || body.Data[0]["owned_by"] != "openai" {
|
||||
t.Errorf("data = %v, want exactly the openai gpt-test entry", body.Data)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGateway_ServeHTTP_GetV1Models_Unauthenticated_Returns401(t *testing.T) {
|
||||
t.Parallel()
|
||||
next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})
|
||||
cfg := CreateConfig()
|
||||
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k", Models: []string{"gpt-test"}}}
|
||||
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
||||
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
|
||||
|
||||
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
|
||||
if rec.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("status = %d, want %d, body: %s", rec.Code, http.StatusUnauthorized, rec.Body.String())
|
||||
}
|
||||
var body struct {
|
||||
Error struct {
|
||||
Type string `json:"type"`
|
||||
} `json:"error"`
|
||||
}
|
||||
if err := json.NewDecoder(rec.Body).Decode(&body); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
if body.Error.Type == "" {
|
||||
t.Error("want a non-empty error.type in the 401 envelope")
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewGateway_WarmFill_PopulatesDiscoveredModels_BeforeAnyRequest proves
|
||||
// newGateway's synchronous first-fill wiring: a discovery-enabled
|
||||
// provider's models are known immediately after New returns, with no
|
||||
// request ever sent.
|
||||
func TestNewGateway_WarmFill_PopulatesDiscoveredModels_BeforeAnyRequest(t *testing.T) {
|
||||
t.Parallel()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"data":[{"id":"disco-model"}]}`))
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})
|
||||
cfg := CreateConfig()
|
||||
cfg.Providers = map[string]*ProviderConfig{
|
||||
"openai": {Type: "openai", APIKey: "k", BaseURL: srv.URL, Discovery: true},
|
||||
}
|
||||
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
gw, ok := h.(*Gateway)
|
||||
if !ok {
|
||||
t.Fatal("handler is not *Gateway")
|
||||
}
|
||||
|
||||
if !gw.registry.states["openai"].hasModel("disco-model") {
|
||||
t.Fatal("want disco-model known immediately after New returns (synchronous warm fill)")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
package traefikllmgateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
)
|
||||
|
||||
// handleModels implements GET /v1/models: an authenticated caller gets back
|
||||
// an OpenAI-compatible {"object":"list","data":[...]} envelope of the
|
||||
// models their group can see, via modelRegistry.listFor.
|
||||
func (g *Gateway) handleModels(w http.ResponseWriter, r *http.Request) {
|
||||
_, grp, ok := g.auth.identify(r)
|
||||
if !ok {
|
||||
writeOAIError(w, http.StatusUnauthorized, "authentication_error", "invalid or missing API key")
|
||||
return
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"object": "list",
|
||||
"data": g.registry.listFor(grp),
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user