diff --git a/llmgateway.go b/llmgateway.go index 46294f5..2fb22f1 100644 --- a/llmgateway.go +++ b/llmgateway.go @@ -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 diff --git a/registry.go b/registry.go new file mode 100644 index 0000000..9b49b20 --- /dev/null +++ b/registry.go @@ -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} +} diff --git a/registry_test.go b/registry_test.go new file mode 100644 index 0000000..848b3a1 --- /dev/null +++ b/registry_test.go @@ -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)") + } +} diff --git a/routes_unified.go b/routes_unified.go new file mode 100644 index 0000000..c503de5 --- /dev/null +++ b/routes_unified.go @@ -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), + }) +}