mirror of
https://github.com/netbirdio/netbird.git
synced 2026-08-24 16:41:30 +02:00
[management] Fetch a provider's model list from the vendor
The catalog is the only source of models an operator can pick from, and it cannot know two things that matter. It goes stale — its entries carry comments recording which models a vendor retired on which date — and it cannot see an account: which OpenAI models an org is entitled to, which Bedrock inference profiles a given account and region hold, which Vertex models a project has enabled. Add a client that asks the vendor directly, with the endpoint, auth header and response shape all declared by the catalog rather than supplied by the caller. The four shapes come from probing the live APIs (see the discovery e2e); each vendor invented its own envelope and none can be guessed from the request. Bedrock is the case that shaped the design. Its listing lives on the control plane while inference must go to the runtime host, so Discovery carries its own host. The ids it returns are region-prefixed and are taken verbatim, because that prefix is what AWS requires and it cannot be derived from the configured region — an eu-central-1 account holds global.* profiles alongside its eu.* ones. The vendor is authoritative for the id; the catalog stays authoritative for pricing. A discovered model the shipped table cannot price is reported as such, so it cannot be registered at a silent zero rate. Management has not made outbound calls on an operator's behalf before, and it holds a credential for every provider, so the host is checked before dialing: every resolved address must be public, which covers the cloud metadata address and NetBird's own overlay range, and redirects are not followed since they would move the request to a host the check never saw.
This commit is contained in:
@@ -113,8 +113,61 @@ type Provider struct {
|
||||
// upstream provider + credentials on Portkey's hosted side).
|
||||
ExtraHeaders []ExtraHeader
|
||||
Models []Model
|
||||
// Discovery, when non-nil, describes how to ask this vendor which
|
||||
// models the operator's own credential can actually reach, so the
|
||||
// provider form can offer a live list instead of only the hand-curated
|
||||
// Models above. Nil for entries with no listing endpoint (gateways
|
||||
// vary too much) — those keep free-text entry.
|
||||
Discovery *Discovery
|
||||
}
|
||||
|
||||
// ListingShape names the response envelope a vendor returns its model
|
||||
// listing in. Every vendor invented its own, and none of them can be
|
||||
// guessed from the request, so the catalog states it.
|
||||
type ListingShape string
|
||||
|
||||
const (
|
||||
// ShapeOpenAIData is {"data":[{"id":…}]} — OpenAI, and Anthropic, which
|
||||
// adopted the same envelope.
|
||||
ShapeOpenAIData ListingShape = "openai_data"
|
||||
// ShapeBedrockInferenceProfiles is
|
||||
// {"inferenceProfileSummaries":[{"inferenceProfileId":…}]}. The ids carry
|
||||
// the region prefix that makes them invocable, which is exactly what an
|
||||
// operator cannot reconstruct by hand.
|
||||
ShapeBedrockInferenceProfiles ListingShape = "bedrock_inference_profiles"
|
||||
// ShapeVertexPublisherModels is {"publisherModels":[{"name":…}]}, where
|
||||
// name is a resource path and the invocable id is its last segment joined
|
||||
// to a separate versionId field.
|
||||
ShapeVertexPublisherModels ListingShape = "vertex_publisher_models"
|
||||
)
|
||||
|
||||
// Discovery describes one vendor's model-listing endpoint.
|
||||
//
|
||||
// Host is deliberately separate from the provider record's upstream URL:
|
||||
// Bedrock serves listings from the control plane (bedrock.<region>) while
|
||||
// inference must go to the runtime host (bedrock-runtime.<region>), so the
|
||||
// two cannot be the same value. Empty Host means "use the record's own
|
||||
// upstream", which is right for every vendor that serves both from one host.
|
||||
//
|
||||
// The regionPlaceholder in Host is substituted from the provider record's
|
||||
// region. Deriving the discovery host from the catalog rather than accepting
|
||||
// one from the caller is also what keeps this from being an open proxy: the
|
||||
// only hosts management will dial are the ones written here.
|
||||
type Discovery struct {
|
||||
Host string
|
||||
Path string
|
||||
Query string
|
||||
Shape ListingShape
|
||||
// Headers are static headers the vendor requires beyond the credential
|
||||
// (Anthropic versions its API through one and rejects a request without
|
||||
// it). The auth header itself comes from AuthHeaderName/Template.
|
||||
Headers map[string]string
|
||||
}
|
||||
|
||||
// RegionPlaceholder is replaced in Discovery.Host by the provider record's
|
||||
// configured region.
|
||||
const RegionPlaceholder = "<region>"
|
||||
|
||||
// ExtraHeader names a single optional per-provider routing/config
|
||||
// header. Catalog declares N of these per provider type; the operator
|
||||
// fills any subset on the provider record (see Provider.ExtraValues).
|
||||
@@ -245,8 +298,12 @@ var providers = []Provider{
|
||||
AuthHeaderTemplate: "Bearer ${API_KEY}",
|
||||
DefaultContentType: "application/json",
|
||||
BrandColor: "#10A37F",
|
||||
ParserID: "openai",
|
||||
PricingSurfaces: []string{"openai"},
|
||||
Discovery: &Discovery{
|
||||
Path: "/v1/models",
|
||||
Shape: ShapeOpenAIData,
|
||||
},
|
||||
ParserID: "openai",
|
||||
PricingSurfaces: []string{"openai"},
|
||||
// Pricing + context windows cross-checked against LiteLLM's
|
||||
// model_prices_and_context_window.json. Notable corrections from
|
||||
// earlier values: o4-mini repriced from $4/$16 to $1.10/$4.40
|
||||
@@ -284,8 +341,18 @@ var providers = []Provider{
|
||||
AuthHeaderTemplate: "${API_KEY}",
|
||||
DefaultContentType: "application/json",
|
||||
BrandColor: "#D97757",
|
||||
ParserID: "anthropic",
|
||||
PricingSurfaces: []string{"anthropic"},
|
||||
Discovery: &Discovery{
|
||||
Path: "/v1/models",
|
||||
// The default page is short and a picker wants the whole
|
||||
// catalogue in one call.
|
||||
Query: "limit=1000",
|
||||
Shape: ShapeOpenAIData,
|
||||
// Anthropic versions its API through a header and refuses a
|
||||
// request that omits it, listing included.
|
||||
Headers: map[string]string{"anthropic-version": "2023-06-01"},
|
||||
},
|
||||
ParserID: "anthropic",
|
||||
PricingSurfaces: []string{"anthropic"},
|
||||
// Per Anthropic's current model lineup. Pricing in USD per 1k
|
||||
// tokens. Context windows: 4.6+ family is 1M; Haiku 4.5 stays at
|
||||
// 200K. claude-3-7-sonnet and claude-3-5-haiku retired
|
||||
@@ -345,6 +412,22 @@ var providers = []Provider{
|
||||
AuthHeaderTemplate: "Bearer ${API_KEY}",
|
||||
DefaultContentType: "application/json",
|
||||
BrandColor: "#FF9900",
|
||||
// Listings come from the CONTROL PLANE, not the runtime host in
|
||||
// DefaultHost above: ListInferenceProfiles is not an operation
|
||||
// bedrock-runtime implements, and answers <UnknownOperationException/>
|
||||
// there. Inference has to go to the runtime host, so the two hosts
|
||||
// genuinely differ and Discovery.Host carries the difference.
|
||||
//
|
||||
// Inference profiles rather than foundation models because the profile
|
||||
// id is the invocable one: it carries the region prefix (eu., us.,
|
||||
// global.) that AWS requires and that cannot be derived from the
|
||||
// configured region — an eu-central-1 account legitimately holds
|
||||
// global.* profiles.
|
||||
Discovery: &Discovery{
|
||||
Host: "bedrock." + RegionPlaceholder + ".amazonaws.com",
|
||||
Path: "/inference-profiles",
|
||||
Shape: ShapeBedrockInferenceProfiles,
|
||||
},
|
||||
// ParserID stays empty (path-style dispatch via IsBedrockPathStyle);
|
||||
// the request parser meters these under the "bedrock" surface.
|
||||
PricingSurfaces: []string{"bedrock"},
|
||||
@@ -395,6 +478,15 @@ var providers = []Provider{
|
||||
AuthHeaderTemplate: "Bearer ${API_KEY}",
|
||||
DefaultContentType: "application/json",
|
||||
BrandColor: "#4285F4",
|
||||
// Only the v1beta1 publisher listing answers: the v1 form and the
|
||||
// project-scoped form under BOTH versions return 404. That means the
|
||||
// list is publisher-global — it cannot say which models this project
|
||||
// has enabled — so it is offered as a suggestion beside the catalog
|
||||
// rather than replacing it. See the discovery e2e for the probes.
|
||||
Discovery: &Discovery{
|
||||
Path: "/v1beta1/publishers/anthropic/models",
|
||||
Shape: ShapeVertexPublisherModels,
|
||||
},
|
||||
// ParserID stays empty (path-style dispatch via IsVertexPathStyle);
|
||||
// Anthropic-on-Vertex requests are metered under the "anthropic"
|
||||
// surface with the bare, unversioned model id.
|
||||
|
||||
@@ -0,0 +1,322 @@
|
||||
// Package modeldiscovery asks a vendor which models an operator's own
|
||||
// credential can reach, so the provider form can offer a live list instead of
|
||||
// only the catalog's hand-curated one.
|
||||
//
|
||||
// The catalog cannot know two things that matter. It goes stale — its entries
|
||||
// carry comments tracking which models a vendor retired on which date — and it
|
||||
// cannot see an account: which OpenAI models an org is entitled to, which
|
||||
// Bedrock inference profiles a given account and region hold, which Vertex
|
||||
// models a project has enabled. Those are exactly the facts an operator needs
|
||||
// when filling in a provider record, and only the vendor has them.
|
||||
//
|
||||
// The vendor is authoritative for the model ID. The catalog remains
|
||||
// authoritative for pricing, and a discovered model the catalog cannot price
|
||||
// is reported as such rather than silently registered at a rate of zero.
|
||||
package modeldiscovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"golang.org/x/oauth2/google"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog"
|
||||
)
|
||||
|
||||
const (
|
||||
// fetchTimeout bounds one vendor call end to end. A listing is a single
|
||||
// small GET; anything slower is a vendor problem and the operator is
|
||||
// waiting on a form.
|
||||
fetchTimeout = 8 * time.Second
|
||||
// maxListingBytes bounds the response we will buffer. The largest real
|
||||
// listing observed is Bedrock's foundation-model catalogue at ~70KB, so
|
||||
// this is a wide margin over anything legitimate.
|
||||
maxListingBytes = 2 << 20
|
||||
// gcpScope matches the scope llm_router mints Vertex tokens under, so a
|
||||
// credential that works for discovery works for inference too.
|
||||
gcpScope = "https://www.googleapis.com/auth/cloud-platform"
|
||||
// vertexKeyfilePrefix marks an api_key that is a base64 service-account
|
||||
// JSON key rather than a bearer token.
|
||||
vertexKeyfilePrefix = "keyfile::"
|
||||
)
|
||||
|
||||
// ErrNoDiscovery is returned for a catalog entry that declares no listing
|
||||
// endpoint. Gateways vary too much to have one, and the caller should fall
|
||||
// back to the catalog list plus free-text entry rather than treating this as
|
||||
// a failure.
|
||||
var ErrNoDiscovery = errors.New("provider has no model-discovery endpoint")
|
||||
|
||||
// Model is one discovered model.
|
||||
type Model struct {
|
||||
// ID is the identifier to register on the provider record, in the form the
|
||||
// vendor issues it. For Bedrock that is the region-prefixed inference
|
||||
// profile id, which is the only form AWS accepts at invoke time.
|
||||
ID string
|
||||
// Label is the vendor's display name where it supplies one.
|
||||
Label string
|
||||
// PricingKnown reports whether the shipped pricing table can price this
|
||||
// model. False means the operator must set rates, or the request would
|
||||
// meter at zero.
|
||||
PricingKnown bool
|
||||
}
|
||||
|
||||
// Request identifies which vendor to ask and with what credential.
|
||||
type Request struct {
|
||||
// CatalogID selects the catalog entry, which supplies the endpoint, the
|
||||
// auth header and the response shape. The caller never supplies those.
|
||||
CatalogID string
|
||||
// UpstreamURL is the provider record's configured upstream. It is used
|
||||
// only when the catalog entry declares no discovery host of its own.
|
||||
UpstreamURL string
|
||||
// Region substitutes the catalog host's <region> placeholder.
|
||||
Region string
|
||||
// APIKey is the operator's credential, exactly as stored on the record.
|
||||
APIKey string
|
||||
}
|
||||
|
||||
// Client fetches model listings. The zero value is usable; Resolver and
|
||||
// HTTPClient exist so tests can drive it against a local server.
|
||||
type Client struct {
|
||||
HTTPClient *http.Client
|
||||
// Resolver looks up the host for the SSRF check. Nil uses the default.
|
||||
Resolver *net.Resolver
|
||||
// AllowPrivateHosts disables the private-address guard. Only tests set it:
|
||||
// their server is on loopback, which is precisely what the guard blocks.
|
||||
AllowPrivateHosts bool
|
||||
}
|
||||
|
||||
// Fetch returns the models the credential can reach.
|
||||
func (c *Client) Fetch(ctx context.Context, req Request) ([]Model, error) {
|
||||
entry, ok := catalog.Lookup(req.CatalogID)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unknown catalog provider %q", req.CatalogID)
|
||||
}
|
||||
if entry.Discovery == nil {
|
||||
return nil, ErrNoDiscovery
|
||||
}
|
||||
|
||||
endpoint, err := c.discoveryURL(entry, req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, fetchTimeout)
|
||||
defer cancel()
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("build discovery request: %w", err)
|
||||
}
|
||||
if err := applyAuth(httpReq, entry, req.APIKey); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for name, value := range entry.Discovery.Headers {
|
||||
httpReq.Header.Set(name, value)
|
||||
}
|
||||
httpReq.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := c.httpClient().Do(httpReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("reach %s: %w", entry.Name, err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, maxListingBytes))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read %s listing: %w", entry.Name, err)
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
// Surface the vendor's own status. An operator whose key lacks a scope
|
||||
// needs to see 403 rather than a generic failure.
|
||||
return nil, fmt.Errorf("%s returned %d for its model listing", entry.Name, resp.StatusCode)
|
||||
}
|
||||
|
||||
ids, err := parseListing(entry.Discovery.Shape, body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return decorate(entry, ids), nil
|
||||
}
|
||||
|
||||
// discoveryURL builds the listing URL and refuses one that does not point at a
|
||||
// public host.
|
||||
//
|
||||
// The path, query and (for Bedrock) the host all come from the catalog rather
|
||||
// than from the caller, so the only operator-controlled part is the host of an
|
||||
// entry whose listing lives on its own upstream. That still has to be checked:
|
||||
// management holds credentials for every provider, and an upstream pointed at
|
||||
// an internal address would turn this endpoint into a probe of the management
|
||||
// server's own network.
|
||||
func (c *Client) discoveryURL(entry catalog.Provider, req Request) (string, error) {
|
||||
host := entry.Discovery.Host
|
||||
if host == "" {
|
||||
parsed, err := url.Parse(strings.TrimSpace(req.UpstreamURL))
|
||||
if err != nil || parsed.Host == "" {
|
||||
return "", fmt.Errorf("provider upstream %q is not a usable URL", req.UpstreamURL)
|
||||
}
|
||||
host = parsed.Host
|
||||
}
|
||||
if strings.Contains(host, catalog.RegionPlaceholder) {
|
||||
region := strings.TrimSpace(req.Region)
|
||||
if region == "" {
|
||||
return "", fmt.Errorf("%s discovery needs a region", entry.Name)
|
||||
}
|
||||
host = strings.ReplaceAll(host, catalog.RegionPlaceholder, region)
|
||||
}
|
||||
|
||||
target := &url.URL{Scheme: "https", Host: host, Path: entry.Discovery.Path, RawQuery: entry.Discovery.Query}
|
||||
if err := c.checkPublicHost(target.Hostname()); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return target.String(), nil
|
||||
}
|
||||
|
||||
// checkPublicHost refuses hosts that resolve to an address the management
|
||||
// server should never be asked to reach on an operator's behalf.
|
||||
func (c *Client) checkPublicHost(host string) error {
|
||||
if c.AllowPrivateHosts {
|
||||
return nil
|
||||
}
|
||||
if host == "" {
|
||||
return errors.New("discovery host is empty")
|
||||
}
|
||||
resolver := c.Resolver
|
||||
if resolver == nil {
|
||||
resolver = net.DefaultResolver
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), fetchTimeout)
|
||||
defer cancel()
|
||||
|
||||
addrs, err := resolver.LookupNetIP(ctx, "ip", host)
|
||||
if err != nil {
|
||||
return fmt.Errorf("resolve discovery host %q: %w", host, err)
|
||||
}
|
||||
// Every address must be public: a name that resolves to one public and one
|
||||
// loopback address is still a way to reach loopback.
|
||||
for _, addr := range addrs {
|
||||
if !isPublic(addr) {
|
||||
return fmt.Errorf("discovery host %q resolves to a non-public address", host)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// isPublic reports whether an address is one we are willing to dial.
|
||||
func isPublic(addr netip.Addr) bool {
|
||||
addr = addr.Unmap()
|
||||
switch {
|
||||
case !addr.IsValid(),
|
||||
addr.IsLoopback(),
|
||||
addr.IsPrivate(),
|
||||
addr.IsLinkLocalUnicast(),
|
||||
addr.IsLinkLocalMulticast(),
|
||||
addr.IsInterfaceLocalMulticast(),
|
||||
addr.IsMulticast(),
|
||||
addr.IsUnspecified():
|
||||
return false
|
||||
}
|
||||
// 100.64.0.0/10 (carrier NAT) is where NetBird's own overlay addresses
|
||||
// live, so it is emphatically not somewhere to send a provider credential.
|
||||
if addr.Is4() {
|
||||
b := addr.As4()
|
||||
if b[0] == 100 && b[1] >= 64 && b[1] <= 127 {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// applyAuth sets the credential header the catalog entry declares. A Vertex
|
||||
// service-account key is exchanged for an OAuth token first, the same way the
|
||||
// proxy does at request time.
|
||||
func applyAuth(req *http.Request, entry catalog.Provider, apiKey string) error {
|
||||
key := strings.TrimSpace(apiKey)
|
||||
if key == "" {
|
||||
return fmt.Errorf("%s discovery needs an API key", entry.Name)
|
||||
}
|
||||
if rest, ok := strings.CutPrefix(key, vertexKeyfilePrefix); ok {
|
||||
token, err := mintGCPToken(req.Context(), rest)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
key = token
|
||||
}
|
||||
name := entry.AuthHeaderName
|
||||
if name == "" {
|
||||
name = "Authorization"
|
||||
}
|
||||
template := entry.AuthHeaderTemplate
|
||||
if template == "" {
|
||||
template = "${API_KEY}"
|
||||
}
|
||||
req.Header.Set(name, strings.ReplaceAll(template, "${API_KEY}", key))
|
||||
return nil
|
||||
}
|
||||
|
||||
// mintGCPToken exchanges a base64 service-account key for an access token.
|
||||
func mintGCPToken(ctx context.Context, saKeyB64 string) (string, error) {
|
||||
jsonKey, err := base64.StdEncoding.DecodeString(strings.TrimSpace(saKeyB64))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("decode service-account key: %w", err)
|
||||
}
|
||||
conf, err := google.JWTConfigFromJSON(jsonKey, gcpScope)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("parse service-account key: %w", err)
|
||||
}
|
||||
tok, err := conf.TokenSource(ctx).Token()
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("mint gcp token: %w", err)
|
||||
}
|
||||
return tok.AccessToken, nil
|
||||
}
|
||||
|
||||
// decorate turns raw vendor ids into the models the caller renders, marking
|
||||
// each with whether the shipped pricing table can price it.
|
||||
func decorate(entry catalog.Provider, ids []listedModel) []Model {
|
||||
priced := make(map[string]struct{}, len(entry.Models))
|
||||
for _, m := range entry.Models {
|
||||
priced[m.ID] = struct{}{}
|
||||
}
|
||||
|
||||
out := make([]Model, 0, len(ids))
|
||||
seen := make(map[string]struct{}, len(ids))
|
||||
for _, listed := range ids {
|
||||
if listed.id == "" {
|
||||
continue
|
||||
}
|
||||
if _, dup := seen[listed.id]; dup {
|
||||
continue
|
||||
}
|
||||
seen[listed.id] = struct{}{}
|
||||
|
||||
// The catalog keys pricing by the normalised id while the vendor
|
||||
// issues the wire form, so normalise before asking whether we can
|
||||
// price it — otherwise every Bedrock profile would report unpriced.
|
||||
_, known := priced[normalizeForPricing(entry.ID, listed.id)]
|
||||
out = append(out, Model{ID: listed.id, Label: listed.label, PricingKnown: known})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (c *Client) httpClient() *http.Client {
|
||||
if c.HTTPClient != nil {
|
||||
return c.HTTPClient
|
||||
}
|
||||
return &http.Client{
|
||||
Timeout: fetchTimeout,
|
||||
// A redirect is a way to move the request to a host the guard above
|
||||
// never checked, so none are followed.
|
||||
CheckRedirect: func(*http.Request, []*http.Request) error {
|
||||
return http.ErrUseLastResponse
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,277 @@
|
||||
package modeldiscovery
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog"
|
||||
)
|
||||
|
||||
// stubTransport answers every request with one canned response and records the
|
||||
// request it was given, so a test can assert on the URL and headers the client
|
||||
// built without a network round trip.
|
||||
type stubTransport struct {
|
||||
status int
|
||||
body string
|
||||
got *http.Request
|
||||
}
|
||||
|
||||
func (s *stubTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
s.got = req
|
||||
status := s.status
|
||||
if status == 0 {
|
||||
status = http.StatusOK
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: status,
|
||||
Body: io.NopCloser(strings.NewReader(s.body)),
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Request: req,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// newStubClient returns a client that never leaves the process. The host guard
|
||||
// is disabled because it would otherwise resolve the vendor's real name, which
|
||||
// would make these tests depend on DNS.
|
||||
func newStubClient(status int, body string) (*Client, *stubTransport) {
|
||||
tr := &stubTransport{status: status, body: body}
|
||||
return &Client{
|
||||
HTTPClient: &http.Client{Transport: tr},
|
||||
AllowPrivateHosts: true,
|
||||
}, tr
|
||||
}
|
||||
|
||||
// The payloads below are trimmed from what the vendors actually returned in
|
||||
// the discovery e2e, rather than invented, so a parser that only works against
|
||||
// an idealised shape fails here.
|
||||
|
||||
const openAIListing = `{"object":"list","data":[
|
||||
{"id":"gpt-4o-mini","object":"model","created":1721172741,"owned_by":"system"},
|
||||
{"id":"gpt-4o","object":"model","created":1715367049,"owned_by":"system"}
|
||||
]}`
|
||||
|
||||
const anthropicListing = `{"data":[
|
||||
{"type":"model","id":"claude-haiku-4-5-20251001","display_name":"Claude Haiku 4.5"},
|
||||
{"type":"model","id":"claude-sonnet-4-6","display_name":"Claude Sonnet 4.6"}
|
||||
],"has_more":false}`
|
||||
|
||||
const bedrockListing = `{"inferenceProfileSummaries":[
|
||||
{"inferenceProfileId":"eu.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"inferenceProfileName":"EU Anthropic Claude Haiku 4.5","status":"ACTIVE","type":"SYSTEM_DEFINED"},
|
||||
{"inferenceProfileId":"global.cohere.embed-v4:0",
|
||||
"inferenceProfileName":"Global Cohere Embed v4","status":"ACTIVE","type":"SYSTEM_DEFINED"},
|
||||
{"inferenceProfileId":"eu.meta.llama3-2-1b-instruct-v1:0",
|
||||
"inferenceProfileName":"EU Meta Llama 3.2 1B","status":"INACTIVE","type":"SYSTEM_DEFINED"}
|
||||
]}`
|
||||
|
||||
const vertexListing = `{"publisherModels":[
|
||||
{"name":"publishers/anthropic/models/claude-3-opus","versionId":"20240229","launchStage":"GA"},
|
||||
{"name":"publishers/anthropic/models/claude-sonnet-4-5","versionId":"20250929","launchStage":"GA"}
|
||||
]}`
|
||||
|
||||
func TestFetchOpenAIListing(t *testing.T) {
|
||||
cl, tr := newStubClient(http.StatusOK, openAIListing)
|
||||
|
||||
models, err := cl.Fetch(context.Background(), Request{
|
||||
CatalogID: "openai_api",
|
||||
UpstreamURL: "https://api.openai.com",
|
||||
APIKey: "sk-test",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, "https://api.openai.com/v1/models", tr.got.URL.String())
|
||||
assert.Equal(t, "Bearer sk-test", tr.got.Header.Get("Authorization"),
|
||||
"the credential must be injected through the catalog's auth template")
|
||||
assert.Equal(t, []string{"gpt-4o-mini", "gpt-4o"}, ids(models))
|
||||
for _, m := range models {
|
||||
assert.True(t, m.PricingKnown, "both models are in the shipped catalog: %s", m.ID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchAnthropicSendsTheVersionHeader(t *testing.T) {
|
||||
cl, tr := newStubClient(http.StatusOK, anthropicListing)
|
||||
|
||||
models, err := cl.Fetch(context.Background(), Request{
|
||||
CatalogID: "anthropic_api",
|
||||
UpstreamURL: "https://api.anthropic.com",
|
||||
APIKey: "sk-ant-test",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Anthropic rejects a request without the version header, so a listing
|
||||
// that reached us at all proves it was sent — but assert it, because the
|
||||
// failure mode otherwise only shows up against the live API.
|
||||
assert.Equal(t, "2023-06-01", tr.got.Header.Get("anthropic-version"))
|
||||
assert.Equal(t, "sk-ant-test", tr.got.Header.Get("x-api-key"),
|
||||
"Anthropic takes a bare key under its own header, not a Bearer token")
|
||||
assert.Equal(t, "limit=1000", tr.got.URL.RawQuery)
|
||||
|
||||
assert.Equal(t, []string{"claude-haiku-4-5-20251001", "claude-sonnet-4-6"}, ids(models))
|
||||
assert.Equal(t, "Claude Haiku 4.5", models[0].Label)
|
||||
}
|
||||
|
||||
func TestFetchBedrockUsesTheControlPlaneAndKeepsWireIDs(t *testing.T) {
|
||||
cl, tr := newStubClient(http.StatusOK, bedrockListing)
|
||||
|
||||
models, err := cl.Fetch(context.Background(), Request{
|
||||
CatalogID: "bedrock_api",
|
||||
// The record's upstream is the RUNTIME host, which does not serve
|
||||
// listings. The catalog's own discovery host must win over it.
|
||||
UpstreamURL: "https://bedrock-runtime.eu-central-1.amazonaws.com",
|
||||
Region: "eu-central-1",
|
||||
APIKey: "aws-bearer",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, "https://bedrock.eu-central-1.amazonaws.com/inference-profiles",
|
||||
tr.got.URL.String(), "listings come from the control plane, not the runtime host")
|
||||
|
||||
// Region-prefixed ids verbatim: the prefix is what makes them invocable
|
||||
// and it cannot be reconstructed — global.* alongside eu.* is exactly the
|
||||
// case that defeats deriving it from the configured region.
|
||||
assert.Equal(t, []string{
|
||||
"eu.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"global.cohere.embed-v4:0",
|
||||
}, ids(models), "an INACTIVE profile must not be offered")
|
||||
|
||||
assert.True(t, models[0].PricingKnown,
|
||||
"the catalog prices anthropic.claude-haiku-4-5, which this id normalises to")
|
||||
assert.False(t, models[1].PricingKnown,
|
||||
"cohere embed is not in the shipped Bedrock catalog, so the operator must price it")
|
||||
}
|
||||
|
||||
func TestFetchVertexJoinsNameAndVersion(t *testing.T) {
|
||||
cl, _ := newStubClient(http.StatusOK, vertexListing)
|
||||
|
||||
models, err := cl.Fetch(context.Background(), Request{
|
||||
CatalogID: "vertex_ai_api",
|
||||
UpstreamURL: "https://us-east5-aiplatform.googleapis.com",
|
||||
Region: "us-east5",
|
||||
APIKey: "ya29.test-token",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Vertex addresses a model as "<id>@<version>" on rawPredict, and splits
|
||||
// those across two fields in the listing.
|
||||
assert.Equal(t, []string{"claude-3-opus@20240229", "claude-sonnet-4-5@20250929"}, ids(models))
|
||||
assert.Equal(t, "claude-3-opus", models[0].Label)
|
||||
}
|
||||
|
||||
func TestFetchSurfacesTheVendorStatus(t *testing.T) {
|
||||
cl, _ := newStubClient(http.StatusForbidden, `{"error":{"message":"no access"}}`)
|
||||
|
||||
_, err := cl.Fetch(context.Background(), Request{
|
||||
CatalogID: "openai_api",
|
||||
UpstreamURL: "https://api.openai.com",
|
||||
APIKey: "sk-test",
|
||||
})
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "403",
|
||||
"an operator whose key lacks access needs to see which status the vendor returned")
|
||||
}
|
||||
|
||||
func TestFetchRejectsAProviderWithoutDiscovery(t *testing.T) {
|
||||
cl, _ := newStubClient(http.StatusOK, openAIListing)
|
||||
|
||||
_, err := cl.Fetch(context.Background(), Request{
|
||||
CatalogID: "litellm_proxy",
|
||||
UpstreamURL: "https://gateway.example.com",
|
||||
APIKey: "sk-test",
|
||||
})
|
||||
assert.ErrorIs(t, err, ErrNoDiscovery,
|
||||
"a gateway with no listing endpoint must be distinguishable from a failure, so the caller can fall back")
|
||||
}
|
||||
|
||||
func TestFetchRequiresACredential(t *testing.T) {
|
||||
cl, _ := newStubClient(http.StatusOK, openAIListing)
|
||||
|
||||
_, err := cl.Fetch(context.Background(), Request{
|
||||
CatalogID: "openai_api",
|
||||
UpstreamURL: "https://api.openai.com",
|
||||
})
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "API key")
|
||||
}
|
||||
|
||||
func TestDiscoveryURLNeedsARegionWhenTheHostTemplatesOne(t *testing.T) {
|
||||
cl, _ := newStubClient(http.StatusOK, bedrockListing)
|
||||
|
||||
_, err := cl.Fetch(context.Background(), Request{
|
||||
CatalogID: "bedrock_api",
|
||||
UpstreamURL: "https://bedrock-runtime.eu-central-1.amazonaws.com",
|
||||
APIKey: "aws-bearer",
|
||||
})
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "region",
|
||||
"an unsubstituted <region> placeholder would dial a host that does not exist")
|
||||
}
|
||||
|
||||
// TestHostGuardRejectsNonPublicAddresses is the SSRF guard. Management holds a
|
||||
// credential for every provider, so an upstream pointed at an internal address
|
||||
// would turn discovery into a way to probe — and hand a token to — the
|
||||
// management server's own network.
|
||||
func TestHostGuardRejectsNonPublicAddresses(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
addr string
|
||||
want bool
|
||||
}{
|
||||
{"loopback v4", "127.0.0.1", false},
|
||||
{"loopback v6", "::1", false},
|
||||
{"private 10/8", "10.0.0.5", false},
|
||||
{"private 172.16/12", "172.16.4.1", false},
|
||||
{"private 192.168/16", "192.168.1.1", false},
|
||||
{"link-local", "169.254.169.254", false}, // cloud metadata
|
||||
{"unspecified", "0.0.0.0", false},
|
||||
{"multicast", "224.0.0.1", false},
|
||||
{"netbird overlay 100.64/10", "100.90.1.2", false},
|
||||
{"v4-mapped loopback", "::ffff:127.0.0.1", false},
|
||||
{"public v4", "1.1.1.1", true},
|
||||
{"public v6", "2606:4700:4700::1111", true},
|
||||
{"just outside CGNAT", "100.128.0.1", true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
addr, err := netip.ParseAddr(tc.addr)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tc.want, isPublic(addr))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHostGuardResolvesAndRejectsLocalhost(t *testing.T) {
|
||||
cl := &Client{}
|
||||
err := cl.checkPublicHost("localhost")
|
||||
require.Error(t, err, "a name resolving to loopback must be refused, not just a literal address")
|
||||
assert.Contains(t, err.Error(), "non-public")
|
||||
}
|
||||
|
||||
// TestEveryDiscoveryEntryHasAParser keeps the catalog and the parser table from
|
||||
// drifting: adding a Discovery block with a shape nothing parses would fail
|
||||
// only at runtime, in front of an operator.
|
||||
func TestEveryDiscoveryEntryHasAParser(t *testing.T) {
|
||||
for _, entry := range catalog.All() {
|
||||
if entry.Discovery == nil {
|
||||
continue
|
||||
}
|
||||
t.Run(entry.ID, func(t *testing.T) {
|
||||
assert.NotEmpty(t, entry.Discovery.Path, "a discovery entry needs a path")
|
||||
_, err := parseListing(entry.Discovery.Shape, []byte(`{}`))
|
||||
assert.NoError(t, err, "shape %q has no parser", entry.Discovery.Shape)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func ids(models []Model) []string {
|
||||
out := make([]string, 0, len(models))
|
||||
for _, m := range models {
|
||||
out = append(out, m.ID)
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
package modeldiscovery
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog"
|
||||
sharedllm "github.com/netbirdio/netbird/shared/llm"
|
||||
)
|
||||
|
||||
// listedModel is one entry lifted out of a vendor listing before the catalog
|
||||
// is consulted about it.
|
||||
type listedModel struct {
|
||||
id string
|
||||
label string
|
||||
}
|
||||
|
||||
// parseListing extracts model ids from a vendor listing. Each vendor invented
|
||||
// its own envelope, and the shape is declared by the catalog rather than
|
||||
// sniffed, so a vendor that changes shape fails loudly instead of silently
|
||||
// returning nothing.
|
||||
func parseListing(shape catalog.ListingShape, body []byte) ([]listedModel, error) {
|
||||
switch shape {
|
||||
case catalog.ShapeOpenAIData:
|
||||
return parseOpenAIData(body)
|
||||
case catalog.ShapeBedrockInferenceProfiles:
|
||||
return parseBedrockInferenceProfiles(body)
|
||||
case catalog.ShapeVertexPublisherModels:
|
||||
return parseVertexPublisherModels(body)
|
||||
default:
|
||||
return nil, fmt.Errorf("no parser for listing shape %q", shape)
|
||||
}
|
||||
}
|
||||
|
||||
// parseOpenAIData reads {"data":[{"id":…}]}, which OpenAI defined and
|
||||
// Anthropic adopted. Anthropic additionally supplies display_name.
|
||||
func parseOpenAIData(body []byte) ([]listedModel, error) {
|
||||
var doc struct {
|
||||
Data []struct {
|
||||
ID string `json:"id"`
|
||||
DisplayName string `json:"display_name"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &doc); err != nil {
|
||||
return nil, fmt.Errorf("decode model listing: %w", err)
|
||||
}
|
||||
out := make([]listedModel, 0, len(doc.Data))
|
||||
for _, entry := range doc.Data {
|
||||
out = append(out, listedModel{id: entry.ID, label: entry.DisplayName})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// parseBedrockInferenceProfiles reads
|
||||
// {"inferenceProfileSummaries":[{"inferenceProfileId":…}]}.
|
||||
//
|
||||
// The profile id is taken verbatim because its region prefix (eu., us.,
|
||||
// global.) is what makes it invocable, and it is not derivable from the
|
||||
// configured region — an account in one region legitimately holds global.*
|
||||
// profiles alongside its regional ones.
|
||||
//
|
||||
// Only ACTIVE profiles are offered: AWS reports others, and registering one
|
||||
// would produce a model that routes inside NetBird and fails at AWS.
|
||||
func parseBedrockInferenceProfiles(body []byte) ([]listedModel, error) {
|
||||
var doc struct {
|
||||
Summaries []struct {
|
||||
ID string `json:"inferenceProfileId"`
|
||||
Name string `json:"inferenceProfileName"`
|
||||
Status string `json:"status"`
|
||||
} `json:"inferenceProfileSummaries"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &doc); err != nil {
|
||||
return nil, fmt.Errorf("decode inference-profile listing: %w", err)
|
||||
}
|
||||
out := make([]listedModel, 0, len(doc.Summaries))
|
||||
for _, entry := range doc.Summaries {
|
||||
if entry.Status != "" && !strings.EqualFold(entry.Status, "ACTIVE") {
|
||||
continue
|
||||
}
|
||||
out = append(out, listedModel{id: entry.ID, label: entry.Name})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// parseVertexPublisherModels reads {"publisherModels":[{"name":…}]}, where
|
||||
// name is a resource path ("publishers/anthropic/models/claude-3-opus") and
|
||||
// the version lives in a separate field.
|
||||
//
|
||||
// Vertex addresses a model as "<id>@<version>" on the rawPredict path, so the
|
||||
// two are joined here: reporting the bare name would hand the operator an id
|
||||
// that looks usable and is not.
|
||||
func parseVertexPublisherModels(body []byte) ([]listedModel, error) {
|
||||
var doc struct {
|
||||
Models []struct {
|
||||
Name string `json:"name"`
|
||||
VersionID string `json:"versionId"`
|
||||
} `json:"publisherModels"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &doc); err != nil {
|
||||
return nil, fmt.Errorf("decode publisher-model listing: %w", err)
|
||||
}
|
||||
out := make([]listedModel, 0, len(doc.Models))
|
||||
for _, entry := range doc.Models {
|
||||
id := entry.Name
|
||||
if slash := strings.LastIndex(id, "/"); slash >= 0 {
|
||||
id = id[slash+1:]
|
||||
}
|
||||
if id == "" {
|
||||
continue
|
||||
}
|
||||
label := id
|
||||
if entry.VersionID != "" {
|
||||
id += "@" + entry.VersionID
|
||||
}
|
||||
out = append(out, listedModel{id: id, label: label})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// normalizeForPricing maps a vendor's wire id onto the key the catalog prices
|
||||
// it under. It mirrors the synthesiser's normalizePricingModelID: the two must
|
||||
// agree, or a model reported here as priced would meter at the default rate
|
||||
// instead of the operator's.
|
||||
func normalizeForPricing(catalogProviderID, modelID string) string {
|
||||
switch {
|
||||
case catalog.IsBedrockPathStyle(catalogProviderID):
|
||||
return sharedllm.NormalizeBedrockModel(modelID)
|
||||
case catalog.IsVertexPathStyle(catalogProviderID):
|
||||
return sharedllm.NormalizeVertexModel(modelID)
|
||||
default:
|
||||
return modelID
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user