mirror of
https://github.com/netbirdio/netbird.git
synced 2026-10-13 08:59:08 +02:00
[management] Expose live model discovery on the provider API
Adds POST /api/agent-network/catalog/providers/models, so the provider form can offer the models an operator's own credential can reach instead of only the compiled-in catalog. A caller names a catalog provider and supplies either the key they are typing (the record does not exist yet) or the id of a saved record whose stored credential should be reused — which lets the dashboard refresh a list without ever holding the key. The two are mutually exclusive: accepting both would run an arbitrary credential under the identity of a record the caller may only be permitted to read. When a record id is given, the catalog id and upstream come from the record too, so the credential cannot be aimed at a different vendor's endpoint. Gated on Create rather than Read. This spends the operator's credential against a third party, which is not something a read-only role should be able to make the server do. A provider with no listing endpoint answers 422 rather than 500: the caller falls back to the catalog's own models on that outcome, so it has to be distinguishable from a failure. The region is read back out of the configured upstream by matching it against the catalog's host template, since a provider record has no region field and the operator already encoded one when they set up inference. An upstream matching no template is refused rather than guessed at — a wrong region would dial another account's endpoint.
This commit is contained in:
@@ -0,0 +1,137 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/modeldiscovery"
|
||||
nbcontext "github.com/netbirdio/netbird/management/server/context"
|
||||
"github.com/netbirdio/netbird/shared/auth"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// discoveryManagerStub records what the handler asked for and returns a canned
|
||||
// answer. The Manager interface is embedded rather than implemented: only the
|
||||
// one method is reachable from this handler, and a call to any other should
|
||||
// fail loudly rather than silently return a zero value.
|
||||
type discoveryManagerStub struct {
|
||||
agentnetwork.Manager
|
||||
|
||||
gotReq modeldiscovery.Request
|
||||
gotRecordID string
|
||||
models []modeldiscovery.Model
|
||||
err error
|
||||
}
|
||||
|
||||
func (s *discoveryManagerStub) DiscoverProviderModels(
|
||||
_ context.Context, _, _ string, req modeldiscovery.Request, recordID string,
|
||||
) ([]modeldiscovery.Model, error) {
|
||||
s.gotReq = req
|
||||
s.gotRecordID = recordID
|
||||
return s.models, s.err
|
||||
}
|
||||
|
||||
// postDiscovery drives the handler with an authenticated request.
|
||||
func postDiscovery(t *testing.T, stub *discoveryManagerStub, body string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
h := &handler{manager: stub}
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/agent-network/catalog/providers/models", strings.NewReader(body))
|
||||
req = req.WithContext(nbcontext.SetUserAuthInContext(req.Context(), auth.UserAuth{
|
||||
AccountId: "acc-1",
|
||||
UserId: "user-1",
|
||||
}))
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
h.discoverProviderModels(rec, req)
|
||||
return rec
|
||||
}
|
||||
|
||||
func TestDiscoverModelsReturnsTheVendorList(t *testing.T) {
|
||||
stub := &discoveryManagerStub{models: []modeldiscovery.Model{
|
||||
{ID: "eu.anthropic.claude-haiku-4-5-20251001-v1:0", Label: "EU Claude Haiku 4.5", PricingKnown: true},
|
||||
{ID: "global.cohere.embed-v4:0", Label: "Global Cohere Embed v4"},
|
||||
}}
|
||||
|
||||
rec := postDiscovery(t, stub, `{
|
||||
"catalog_provider_id":"bedrock_api",
|
||||
"upstream_url":"https://bedrock-runtime.eu-central-1.amazonaws.com",
|
||||
"api_key":"aws-bearer"
|
||||
}`)
|
||||
require.Equal(t, http.StatusOK, rec.Code, "body: %s", rec.Body.String())
|
||||
|
||||
var out api.AgentNetworkModelDiscoveryResponse
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &out))
|
||||
require.Len(t, out.Models, 2)
|
||||
|
||||
assert.Equal(t, "eu.anthropic.claude-haiku-4-5-20251001-v1:0", out.Models[0].Id)
|
||||
assert.True(t, out.Models[0].PricingKnown)
|
||||
// An unpriced model must say so rather than arriving indistinguishable
|
||||
// from a priced one: registering it silently would meter at zero.
|
||||
assert.False(t, out.Models[1].PricingKnown)
|
||||
|
||||
assert.Equal(t, "bedrock_api", stub.gotReq.CatalogID)
|
||||
assert.Equal(t, "aws-bearer", stub.gotReq.APIKey)
|
||||
assert.Empty(t, stub.gotRecordID)
|
||||
}
|
||||
|
||||
func TestDiscoverModelsUsesAStoredRecordWithoutAKey(t *testing.T) {
|
||||
stub := &discoveryManagerStub{}
|
||||
|
||||
rec := postDiscovery(t, stub, `{"catalog_provider_id":"openai_api","provider_id":"prov-42"}`)
|
||||
require.Equal(t, http.StatusOK, rec.Code, "body: %s", rec.Body.String())
|
||||
|
||||
// The dashboard refreshes a saved provider's list without ever holding
|
||||
// the credential, so the record id has to reach the manager.
|
||||
assert.Equal(t, "prov-42", stub.gotRecordID)
|
||||
assert.Empty(t, stub.gotReq.APIKey)
|
||||
}
|
||||
|
||||
// TestDiscoverModelsRefusesMixedCredentials covers the case where a caller
|
||||
// names a saved provider AND supplies a key. Accepting it would run an
|
||||
// arbitrary credential under the identity of a record the caller may only be
|
||||
// permitted to read.
|
||||
func TestDiscoverModelsRefusesMixedCredentials(t *testing.T) {
|
||||
stub := &discoveryManagerStub{}
|
||||
|
||||
rec := postDiscovery(t, stub, `{
|
||||
"catalog_provider_id":"openai_api",
|
||||
"provider_id":"prov-42",
|
||||
"api_key":"sk-attacker"
|
||||
}`)
|
||||
assert.Equal(t, http.StatusBadRequest, rec.Code)
|
||||
assert.Empty(t, stub.gotRecordID, "the request must be refused before it reaches the manager")
|
||||
}
|
||||
|
||||
// TestDiscoverModelsReportsNoDiscoveryDistinctly matters because the caller
|
||||
// falls back to the catalog's own model list on this outcome. Collapsing it
|
||||
// into a generic 500 would turn "this provider has no listing endpoint" into
|
||||
// "something went wrong", and the form would show an error instead of a list.
|
||||
func TestDiscoverModelsReportsNoDiscoveryDistinctly(t *testing.T) {
|
||||
stub := &discoveryManagerStub{err: modeldiscovery.ErrNoDiscovery}
|
||||
|
||||
rec := postDiscovery(t, stub, `{"catalog_provider_id":"litellm_proxy","upstream_url":"https://gw.example.com","api_key":"sk"}`)
|
||||
assert.Equal(t, http.StatusUnprocessableEntity, rec.Code)
|
||||
}
|
||||
|
||||
func TestDiscoverModelsRejectsMalformedRequests(t *testing.T) {
|
||||
for name, body := range map[string]string{
|
||||
"not json": `{`,
|
||||
"no catalog provider": `{"api_key":"sk"}`,
|
||||
"blank catalog provide": `{"catalog_provider_id":" ","api_key":"sk"}`,
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
stub := &discoveryManagerStub{}
|
||||
rec := postDiscovery(t, stub, body)
|
||||
assert.Equal(t, http.StatusBadRequest, rec.Code)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -7,6 +7,7 @@ package handlers
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"math"
|
||||
"net/http"
|
||||
"net/url"
|
||||
@@ -16,6 +17,7 @@ import (
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/modeldiscovery"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/pricing"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
nbcontext "github.com/netbirdio/netbird/management/server/context"
|
||||
@@ -32,6 +34,7 @@ type handler struct {
|
||||
func RegisterEndpoints(manager agentnetwork.Manager, router *mux.Router) {
|
||||
h := &handler{manager: manager}
|
||||
router.HandleFunc("/agent-network/catalog/providers", h.getCatalogProviders).Methods("GET", "OPTIONS")
|
||||
router.HandleFunc("/agent-network/catalog/providers/models", h.discoverProviderModels).Methods("POST", "OPTIONS")
|
||||
router.HandleFunc("/agent-network/providers", h.getAllProviders).Methods("GET", "OPTIONS")
|
||||
router.HandleFunc("/agent-network/providers", h.createProvider).Methods("POST", "OPTIONS")
|
||||
router.HandleFunc("/agent-network/providers/{providerId}", h.getProvider).Methods("GET", "OPTIONS")
|
||||
@@ -61,6 +64,73 @@ func (h *handler) getCatalogProviders(w http.ResponseWriter, r *http.Request) {
|
||||
util.WriteJSONObject(r.Context(), w, out)
|
||||
}
|
||||
|
||||
// discoverProviderModels asks the vendor which models the operator's own
|
||||
// credential can reach, so the provider form can offer a live list rather than
|
||||
// only the static catalog.
|
||||
func (h *handler) discoverProviderModels(w http.ResponseWriter, r *http.Request) {
|
||||
userAuth, err := nbcontext.GetUserAuthFromContext(r.Context())
|
||||
if err != nil {
|
||||
util.WriteError(r.Context(), err, w)
|
||||
return
|
||||
}
|
||||
|
||||
var body api.AgentNetworkModelDiscoveryRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
util.WriteErrorResponse("invalid json", http.StatusBadRequest, w)
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(body.CatalogProviderId) == "" {
|
||||
util.WriteErrorResponse("catalog_provider_id is required", http.StatusBadRequest, w)
|
||||
return
|
||||
}
|
||||
|
||||
recordID := strValue(body.ProviderId)
|
||||
req := modeldiscovery.Request{
|
||||
CatalogID: body.CatalogProviderId,
|
||||
UpstreamURL: strValue(body.UpstreamUrl),
|
||||
APIKey: strValue(body.ApiKey),
|
||||
}
|
||||
// One source of credential or the other, never a mix: taking a key from
|
||||
// the request while addressing a saved record would let a caller run an
|
||||
// arbitrary credential against a provider they can only read.
|
||||
if recordID != "" && req.APIKey != "" {
|
||||
util.WriteErrorResponse("provide either provider_id or api_key, not both", http.StatusBadRequest, w)
|
||||
return
|
||||
}
|
||||
|
||||
models, err := h.manager.DiscoverProviderModels(r.Context(), userAuth.AccountId, userAuth.UserId, req, recordID)
|
||||
if err != nil {
|
||||
// A provider with no listing endpoint is a fact about the catalog
|
||||
// entry, not a failure: the caller falls back to the catalog's own
|
||||
// models, so it must be able to tell the two apart.
|
||||
if errors.Is(err, modeldiscovery.ErrNoDiscovery) {
|
||||
util.WriteErrorResponse(err.Error(), http.StatusUnprocessableEntity, w)
|
||||
return
|
||||
}
|
||||
util.WriteError(r.Context(), err, w)
|
||||
return
|
||||
}
|
||||
|
||||
out := api.AgentNetworkModelDiscoveryResponse{Models: make([]api.AgentNetworkDiscoveredModel, 0, len(models))}
|
||||
for _, m := range models {
|
||||
entry := api.AgentNetworkDiscoveredModel{Id: m.ID, PricingKnown: m.PricingKnown}
|
||||
if m.Label != "" {
|
||||
label := m.Label
|
||||
entry.Label = &label
|
||||
}
|
||||
out.Models = append(out.Models, entry)
|
||||
}
|
||||
util.WriteJSONObject(r.Context(), w, out)
|
||||
}
|
||||
|
||||
// strValue reads an optional string field, treating absent as empty.
|
||||
func strValue(v *string) string {
|
||||
if v == nil {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(*v)
|
||||
}
|
||||
|
||||
// applyDefaultPricing overwrites the catalog response's model rates with
|
||||
// the LIVE default pricing table, which may differ from the compiled-in
|
||||
// catalog rates when the operator provides a defaults_llm_pricing.yaml.
|
||||
|
||||
Reference in New Issue
Block a user