Compare commits

..

5 Commits

Author SHA1 Message Date
Brandon Hopkins
fea131101d API-key normalization and model-discovery snapshot validation 2026-07-26 18:58:39 -07:00
Brandon Hopkins
0463b30152 Ollama Cloud provider integration and backend test coverage 2026-07-26 18:19:38 -07:00
Brandon Hopkins
d3909e4faf endpoint model discovery and proxy integration 2026-07-26 16:38:01 -07:00
Brandon Hopkins
a92cdb7dcd authentication semantics, API contracts, and tests 2026-07-26 15:37:41 -07:00
Brandon Hopkins
564abb660c Add self-hosted Ollama 2026-07-26 15:05:01 -07:00
53 changed files with 4295 additions and 708 deletions

View File

@@ -273,8 +273,8 @@ dockers_v2:
- netbirdio/netbird
- ghcr.io/netbirdio/netbird
tags:
- "{{ .Version }}-rootless"
- "{{ if eq .Env.SKIP_PUBLISH \"false\" }}rootless-latest{{ end }}"
- "v{{ .Version }}-rootless"
- "{{ if eq .Env.SKIP_PUBLISH \"false\" }}latest{{ end }}"
dockerfile: client/Dockerfile-rootless
extra_files:
- client/netbird-entrypoint.sh

View File

@@ -21,8 +21,8 @@ import (
"github.com/netbirdio/netbird/client/internal"
"github.com/netbirdio/netbird/client/internal/peer"
"github.com/netbirdio/netbird/client/internal/profilemanager"
nbnet "github.com/netbirdio/netbird/client/net"
"github.com/netbirdio/netbird/client/proto"
nbnet "github.com/netbirdio/netbird/client/net"
"github.com/netbirdio/netbird/client/server"
"github.com/netbirdio/netbird/client/system"
"github.com/netbirdio/netbird/shared/management/domain"

View File

@@ -702,6 +702,7 @@ func ifname(n string) []byte {
return b
}
// ipToBytes converts net.IP to the correct byte length for the address family.
func ipToBytes(ip net.IP, af addrFamily) []byte {
if af.addrLen == 4 {
@@ -709,3 +710,4 @@ func ipToBytes(ip net.IP, af addrFamily) []byte {
}
return ip.To16()
}

View File

@@ -158,8 +158,8 @@ func TestTCPPerStateTeardownTimeouts(t *testing.T) {
name: "FinWait2",
drive: func(t *testing.T, tr *TCPTracker, srcIP netip.Addr, srcPort uint16) (ConnKey, TCPState) {
establishConnection(t, tr, srcIP, dstIP, srcPort, dstPort)
tr.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPFin|TCPAck, 0) // FinWait1
require.True(t, tr.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPAck, 0)) // → FinWait2
tr.TrackOutbound(srcIP, dstIP, srcPort, dstPort, TCPFin|TCPAck, 0) // FinWait1
require.True(t, tr.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPAck, 0)) // → FinWait2
return ConnKey{SrcIP: srcIP, DstIP: dstIP, SrcPort: srcPort, DstPort: dstPort}, TCPStateFinWait2
},
},

View File

@@ -18,7 +18,7 @@ type PeerRule struct {
protoLayer gopacket.LayerType
sPort *firewall.Port
dPort *firewall.Port
drop bool
drop bool
}
// ID returns the rule id

View File

@@ -34,14 +34,14 @@ func (f *fakeTunDevice) Create() (device.WGConfigurer, error) {
func (f *fakeTunDevice) Up() (*udpmux.UniversalUDPMuxDefault, error) {
return nil, errors.New("not implemented")
}
func (f *fakeTunDevice) UpdateAddr(wgaddr.Address) error { return nil }
func (f *fakeTunDevice) WgAddress() wgaddr.Address { return wgaddr.Address{} }
func (f *fakeTunDevice) MTU() uint16 { return DefaultMTU }
func (f *fakeTunDevice) DeviceName() string { return "nb-close-test" }
func (f *fakeTunDevice) UpdateAddr(wgaddr.Address) error { return nil }
func (f *fakeTunDevice) WgAddress() wgaddr.Address { return wgaddr.Address{} }
func (f *fakeTunDevice) MTU() uint16 { return DefaultMTU }
func (f *fakeTunDevice) DeviceName() string { return "nb-close-test" }
func (f *fakeTunDevice) FilteredDevice() *device.FilteredDevice { return nil }
func (f *fakeTunDevice) Device() *wgdevice.Device { return nil }
func (f *fakeTunDevice) GetNet() *netstack.Net { return nil }
func (f *fakeTunDevice) GetICEBind() device.EndpointManager { return nil }
func (f *fakeTunDevice) Device() *wgdevice.Device { return nil }
func (f *fakeTunDevice) GetNet() *netstack.Net { return nil }
func (f *fakeTunDevice) GetICEBind() device.EndpointManager { return nil }
func (f *fakeTunDevice) Close() error {
close(f.closeStarted)

View File

@@ -116,11 +116,11 @@ func (d *DefaultManager) ApplyFiltering(networkMap *mgmProto.NetworkMap, dnsRout
// firewall state, so an identical hash means an identical resulting ruleset.
func (d *DefaultManager) firewallConfigHash(networkMap *mgmProto.NetworkMap, dnsRouteFeatureFlag bool) (uint64, error) {
return hashstructure.Hash(struct {
PeerRules []*mgmProto.FirewallRule
PeerRulesIsEmpty bool
RouteRules []*mgmProto.RouteFirewallRule
RouteRulesIsEmpty bool
DNSRouteFeatureFlag bool
PeerRules []*mgmProto.FirewallRule
PeerRulesIsEmpty bool
RouteRules []*mgmProto.RouteFirewallRule
RouteRulesIsEmpty bool
DNSRouteFeatureFlag bool
}{
PeerRules: networkMap.GetFirewallRules(),
PeerRulesIsEmpty: networkMap.GetFirewallRulesIsEmpty(),
@@ -407,6 +407,7 @@ func (d *DefaultManager) getRuleGroupingSelector(rule *mgmProto.FirewallRule) st
return fmt.Sprintf("%v:%v:%v:%s:%v", strconv.Itoa(int(rule.Direction)), rule.Action, rule.Protocol, rule.Port, rule.PortInfo)
}
// extractRuleIP extracts the peer IP from a firewall rule.
// If sourcePrefixes is populated (new management), decode the first entry and use its address.
// Otherwise fall back to the deprecated PeerIP string field (old management).

View File

@@ -556,7 +556,7 @@ func TestApplyFilteringSkipsUnchangedConfig(t *testing.T) {
func buildNetworkMap(peerRules, routeRules int) *mgmProto.NetworkMap {
nm := &mgmProto.NetworkMap{
FirewallRulesIsEmpty: peerRules == 0,
FirewallRulesIsEmpty: peerRules == 0,
RoutesFirewallRulesIsEmpty: routeRules == 0,
}
for i := range peerRules {

View File

@@ -49,3 +49,4 @@ func (tf *GeneralManager) FreeDNSFwd() error {
log.Debugf("free ebpf DNS forwarder")
return tf.unsetFeatureFlag(featureFlagDnsForwarder)
}

View File

@@ -138,6 +138,7 @@ func TestManager_Cleanup_NilMapping(t *testing.T) {
m.cleanup(context.Background(), gateway)
}
func TestManager_CreateMapping_PermanentLeaseFallback(t *testing.T) {
m := NewManager()
m.wgPort = 51820

View File

@@ -1,12 +0,0 @@
//go:build ios
package NetBirdSDK
import "github.com/netbirdio/netbird/version"
// GoClientVersion returns the NetBird Go client version that was baked into
// the framework at compile time via
// -ldflags "-X github.com/netbirdio/netbird/version.version=<version>".
func GoClientVersion() string {
return version.NetbirdVersion()
}

View File

@@ -193,3 +193,4 @@ func buildAddressList(hostname string, remote net.Addr) []string {
}
return addresses
}

View File

@@ -11,7 +11,6 @@ package main
// - The in-process StatusNotifierWatcher + XEmbed host used on minimal WMs
// (Fluxbox, i3, dwm, OpenBox) maps a Button1 press to that same Activate
// call itself (xembed_host_linux.go), so it routes through the same hook.
//
// Registering OnClick here therefore covers both paths with one handler — no
// changes to the watcher or XEmbed C code are needed. Left-click now opens the
// main window; right-click still opens the menu via Wails' default

View File

@@ -72,12 +72,12 @@ type proxyConnection struct {
// Go-side callbacks exposed to JS. js.FuncOf pins the Go closure in a
// global handle map and MUST be released, otherwise every connection
// leaks the Go memory the closure captures.
wsHandlerFn js.Func
onMessageFn js.Func
onCloseFn js.Func
cleanupOnce sync.Once
ctx context.Context
cancel context.CancelFunc
wsHandlerFn js.Func
onMessageFn js.Func
onCloseFn js.Func
cleanupOnce sync.Once
ctx context.Context
cancel context.CancelFunc
}
// NewRDCleanPathProxy creates a new RDCleanPath proxy

View File

@@ -3,9 +3,9 @@ package controller
import (
"context"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
"github.com/netbirdio/netbird/management/internals/modules/zones"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
"github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/store"
"github.com/netbirdio/netbird/management/server/types"

View File

@@ -25,8 +25,8 @@ type Model struct {
// These typically need NetBird identity stamped onto upstream
// requests so the gateway's analytics and budgets attribute to the
// real caller; that's what IdentityInjection is for.
// - KindCustom: the catch-all "OpenAI-compatible self-hosted endpoint"
// entry (vLLM, Ollama, custom inference servers).
// - KindCustom: named and generic OpenAI-compatible self-hosted
// endpoints (vLLM, Ollama, custom inference servers).
//
// Frontend uses Kind to group the provider Select in the modal so an
// operator can spot at a glance which catalog entries proxy other
@@ -40,6 +40,18 @@ const (
KindCustom ProviderKind = "custom"
)
// AuthMode describes whether a catalog provider accepts an upstream
// credential. The zero value intentionally resolves to AuthModeRequired so
// existing and newly-added catalog entries fail closed unless they explicitly
// opt into a less restrictive mode.
type AuthMode string
const (
AuthModeRequired AuthMode = "required"
AuthModeOptional AuthMode = "optional"
AuthModeNone AuthMode = "none"
)
// Provider is the in-memory representation of a catalog provider.
type Provider struct {
ID string
@@ -48,6 +60,9 @@ type Provider struct {
DefaultHost string
// Kind groups this entry for UI presentation; see ProviderKind.
Kind ProviderKind
// AuthMode declares whether the provider requires, optionally accepts, or
// does not support an upstream API key. Empty defaults to required.
AuthMode AuthMode
// AuthHeaderName is the HTTP header the provider's API expects
// the credential under (e.g. "Authorization" for OpenAI,
// "x-api-key" for Anthropic). Combined with AuthHeaderTemplate
@@ -86,6 +101,17 @@ type Provider struct {
// upstream provider + credentials on Portkey's hosted side).
ExtraHeaders []ExtraHeader
Models []Model
// ModelDiscovery configures provider-specific discovery behavior. A nil
// profile means discovery is unsupported. The profile never carries
// caller-supplied paths: the proxy owns the fixed endpoint allowlist.
ModelDiscovery *ModelDiscovery
}
// ModelDiscovery describes the safe, catalog-owned fallback behavior for a
// discoverable provider. Every discovery starts with OpenAI-compatible
// /v1/models; OllamaFallback permits /api/tags only when that route is absent.
type ModelDiscovery struct {
OllamaFallback bool
}
// ExtraHeader names a single optional per-provider routing/config
@@ -461,6 +487,27 @@ var providers = []Provider{
{ID: "kimi-k3", Label: "Kimi K3", InputPer1k: 0.003, OutputPer1k: 0.015, ContextWindow: 1000000},
},
},
{
// Direct Ollama Cloud uses the hosted OpenAI-compatible /v1 API.
// Models are discovered dynamically because the Cloud catalog changes
// independently of NetBird releases. ParserID intentionally remains
// empty to preserve the routing behavior of Ollama, vLLM, and custom
// OpenAI-compatible providers.
ID: "ollama_cloud",
Kind: KindProvider,
AuthMode: AuthModeRequired,
Name: "Ollama Cloud",
Description: "Hosted Ollama models via the OpenAI-compatible API",
DefaultHost: "ollama.com",
AuthHeaderName: "Authorization",
AuthHeaderTemplate: "Bearer ${API_KEY}",
DefaultContentType: "application/json",
BrandColor: "#000000",
Models: []Model{},
ModelDiscovery: &ModelDiscovery{
OllamaFallback: true,
},
},
{
ID: "litellm_proxy",
Kind: KindGateway,
@@ -701,11 +748,31 @@ var providers = []Provider{
BrandColor: "#30A2FF",
Models: []Model{},
},
{
// Ollama exposes an OpenAI-compatible /v1 API. Like vLLM, it gets a
// dedicated catalog id for provider-specific setup guidance while
// retaining generic custom-provider routing. Authentication is optional
// because local Ollama is authless but protected front ends may use it.
ID: "ollama",
Kind: KindCustom,
AuthMode: AuthModeOptional,
Name: "Ollama",
Description: "Self-hosted Ollama (OpenAI-compatible)",
DefaultHost: "",
AuthHeaderName: "Authorization",
AuthHeaderTemplate: "Bearer ${API_KEY}",
DefaultContentType: "application/json",
BrandColor: "#000000",
Models: []Model{},
ModelDiscovery: &ModelDiscovery{
OllamaFallback: true,
},
},
{
ID: "custom",
Kind: KindCustom,
Name: "Custom / Self-hosted",
Description: "OpenAI-compatible endpoint (vLLM, Ollama, …)",
Description: "Other OpenAI-compatible endpoint",
DefaultHost: "",
AuthHeaderName: "Authorization",
AuthHeaderTemplate: "Bearer ${API_KEY}",
@@ -738,6 +805,15 @@ func IsKnown(id string) bool {
return ok
}
// EffectiveAuthMode returns the provider's configured authentication mode.
// Treating the zero value as required keeps older catalog entries fail closed.
func (p Provider) EffectiveAuthMode() AuthMode {
if p.AuthMode == "" {
return AuthModeRequired
}
return p.AuthMode
}
// IsVertexPathStyle reports whether a provider uses the Google Vertex AI
// request shape — the model is carried in the URL path
// (/v1/projects/{p}/locations/{r}/publishers/{pub}/models/{model}:{action})
@@ -774,15 +850,17 @@ func (p Provider) ToAPIResponse() api.AgentNetworkCatalogProvider {
kind = api.AgentNetworkCatalogProviderKindCustom
}
resp := api.AgentNetworkCatalogProvider{
Id: p.ID,
Name: p.Name,
Description: p.Description,
DefaultHost: p.DefaultHost,
Kind: kind,
AuthHeaderTemplate: p.AuthHeaderTemplate,
DefaultContentType: p.DefaultContentType,
BrandColor: p.BrandColor,
Models: models,
Id: p.ID,
Name: p.Name,
Description: p.Description,
DefaultHost: p.DefaultHost,
Kind: kind,
AuthMode: api.AgentNetworkCatalogProviderAuthMode(p.EffectiveAuthMode()),
AuthHeaderTemplate: p.AuthHeaderTemplate,
DefaultContentType: p.DefaultContentType,
BrandColor: p.BrandColor,
Models: models,
SupportsModelDiscovery: p.ModelDiscovery != nil,
}
if len(p.ExtraHeaders) > 0 {
extras := make([]api.AgentNetworkCatalogExtraHeader, 0, len(p.ExtraHeaders))

View File

@@ -0,0 +1,96 @@
package catalog
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/shared/management/http/api"
)
func TestOllamaCatalogEntry(t *testing.T) {
entry, ok := Lookup("ollama")
require.True(t, ok, "Ollama must be available as a dedicated catalog provider")
assert.Equal(t, KindCustom, entry.Kind)
assert.Equal(t, AuthModeOptional, entry.EffectiveAuthMode())
assert.Equal(t, "Ollama", entry.Name)
assert.Equal(t, "Self-hosted Ollama (OpenAI-compatible)", entry.Description)
assert.Empty(t, entry.DefaultHost, "the dashboard owns Ollama's scheme-aware HTTP placeholder")
assert.Equal(t, "Authorization", entry.AuthHeaderName)
assert.Equal(t, "Bearer ${API_KEY}", entry.AuthHeaderTemplate)
assert.Equal(t, "application/json", entry.DefaultContentType)
assert.Empty(t, entry.ParserID, "Ollama preserves the untagged vLLM/custom routing behavior")
assert.Empty(t, entry.Models, "Ollama models are installed dynamically on the configured endpoint")
require.NotNil(t, entry.ModelDiscovery)
assert.True(t, entry.ModelDiscovery.OllamaFallback)
wire := entry.ToAPIResponse()
assert.Equal(t, "ollama", wire.Id)
assert.Equal(t, api.AgentNetworkCatalogProviderKindCustom, wire.Kind)
assert.Equal(t, api.AgentNetworkCatalogProviderAuthModeOptional, wire.AuthMode)
assert.True(t, wire.SupportsModelDiscovery)
assert.NotNil(t, wire.Models)
assert.Empty(t, wire.Models)
}
func TestOllamaCloudCatalogEntry(t *testing.T) {
entry, ok := Lookup("ollama_cloud")
require.True(t, ok, "Ollama Cloud must be available as a dedicated catalog provider")
assert.Equal(t, KindProvider, entry.Kind)
assert.Equal(t, AuthModeRequired, entry.EffectiveAuthMode())
assert.Equal(t, "Ollama Cloud", entry.Name)
assert.Equal(t, "Hosted Ollama models via the OpenAI-compatible API", entry.Description)
assert.Equal(t, "ollama.com", entry.DefaultHost)
assert.Equal(t, "Authorization", entry.AuthHeaderName)
assert.Equal(t, "Bearer ${API_KEY}", entry.AuthHeaderTemplate)
assert.Equal(t, "application/json", entry.DefaultContentType)
assert.Empty(t, entry.ParserID, "Ollama Cloud preserves the untagged Ollama/vLLM/custom routing behavior")
assert.Empty(t, entry.Models, "Ollama Cloud models are discovered dynamically")
require.NotNil(t, entry.ModelDiscovery)
assert.True(t, entry.ModelDiscovery.OllamaFallback)
wire := entry.ToAPIResponse()
assert.Equal(t, "ollama_cloud", wire.Id)
assert.Equal(t, api.AgentNetworkCatalogProviderKindProvider, wire.Kind)
assert.Equal(t, api.AgentNetworkCatalogProviderAuthModeRequired, wire.AuthMode)
assert.Equal(t, "ollama.com", wire.DefaultHost)
assert.True(t, wire.SupportsModelDiscovery)
assert.NotNil(t, wire.Models)
assert.Empty(t, wire.Models)
}
func TestOnlyOllamaProvidersSupportModelDiscovery(t *testing.T) {
discoverable := map[string]bool{
"ollama": true,
"ollama_cloud": true,
}
for _, entry := range All() {
supportsDiscovery := entry.ModelDiscovery != nil
assert.Equal(t, discoverable[entry.ID], supportsDiscovery, entry.ID)
assert.Equal(t, supportsDiscovery, entry.ToAPIResponse().SupportsModelDiscovery, entry.ID)
}
}
func TestCatalogAuthenticationModes(t *testing.T) {
openAI, ok := Lookup("openai_api")
require.True(t, ok)
assert.Empty(t, openAI.AuthMode, "existing entries use the fail-closed default")
assert.Equal(t, AuthModeRequired, openAI.EffectiveAuthMode())
assert.Equal(t, api.AgentNetworkCatalogProviderAuthModeRequired, openAI.ToAPIResponse().AuthMode)
for _, entry := range All() {
switch entry.EffectiveAuthMode() {
case AuthModeRequired, AuthModeOptional:
assert.NotEmpty(t, entry.AuthHeaderName, "%s must declare an auth header name", entry.ID)
assert.NotEmpty(t, entry.AuthHeaderTemplate, "%s must declare an auth header template", entry.ID)
case AuthModeNone:
assert.Empty(t, entry.AuthHeaderName, "%s must not declare an auth header name", entry.ID)
assert.Empty(t, entry.AuthHeaderTemplate, "%s must not declare an auth header template", entry.ID)
default:
t.Errorf("%s has invalid auth mode %q", entry.ID, entry.AuthMode)
}
}
}

View File

@@ -0,0 +1,72 @@
package handlers
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"github.com/gorilla/mux"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
nbcontext "github.com/netbirdio/netbird/management/server/context"
"github.com/netbirdio/netbird/shared/auth"
"github.com/netbirdio/netbird/shared/management/http/api"
)
type discoveryManagerStub struct {
agentnetwork.Manager
result *agentNetworkTypes.ModelDiscoveryResult
err error
accountID string
userID string
providerID string
}
func (m *discoveryManagerStub) DiscoverProviderModels(_ context.Context, accountID, userID, providerID string) (*agentNetworkTypes.ModelDiscoveryResult, error) {
m.accountID = accountID
m.userID = userID
m.providerID = providerID
return m.result, m.err
}
func TestDiscoverProviderModelsHandler(t *testing.T) {
manager := &discoveryManagerStub{
Manager: agentnetwork.NewManagerMock(),
result: &agentNetworkTypes.ModelDiscoveryResult{
RequestID: "probe-123",
Source: "ollama_api_tags",
ProxyCluster: "private.example.com",
Models: []agentNetworkTypes.DiscoveredModel{
{ID: "llama3.2:latest", Label: "llama3.2:latest"},
},
},
}
router := mux.NewRouter()
RegisterEndpoints(manager, router)
req := httptest.NewRequest(http.MethodPost, "/agent-network/providers/provider-1/discover-models", nil)
req = nbcontext.SetUserAuthInRequest(req, auth.UserAuth{
UserId: testUserID,
AccountId: testAccountID,
})
rec := httptest.NewRecorder()
router.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
assert.Equal(t, "no-store", rec.Header().Get("Cache-Control"))
var response api.AgentNetworkModelDiscoveryResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &response))
assert.Equal(t, testAccountID, manager.accountID)
assert.Equal(t, testUserID, manager.userID)
assert.Equal(t, "provider-1", manager.providerID)
assert.Equal(t, "probe-123", response.RequestId)
assert.Equal(t, "private.example.com", response.ProxyCluster)
assert.Equal(t, api.AgentNetworkModelDiscoveryResponseSourceOllamaApiTags, response.Source)
require.Len(t, response.Models, 1)
assert.Equal(t, "llama3.2:latest", response.Models[0].Id)
}

View File

@@ -35,6 +35,7 @@ func RegisterEndpoints(manager agentnetwork.Manager, router *mux.Router) {
router.HandleFunc("/agent-network/providers/{providerId}", h.getProvider).Methods("GET", "OPTIONS")
router.HandleFunc("/agent-network/providers/{providerId}", h.updateProvider).Methods("PUT", "OPTIONS")
router.HandleFunc("/agent-network/providers/{providerId}", h.deleteProvider).Methods("DELETE", "OPTIONS")
router.HandleFunc("/agent-network/providers/{providerId}/discover-models", h.discoverProviderModels).Methods("POST", "OPTIONS")
h.addPolicyEndpoints(router)
h.addGuardrailEndpoints(router)
h.addSettingsEndpoints(router)
@@ -98,6 +99,41 @@ func (h *handler) getProvider(w http.ResponseWriter, r *http.Request) {
util.WriteJSONObject(r.Context(), w, provider.ToAPIResponse())
}
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
}
providerID := strings.TrimSpace(mux.Vars(r)["providerId"])
if providerID == "" {
util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "provider ID is required"), w)
return
}
result, err := h.manager.DiscoverProviderModels(r.Context(), userAuth.AccountId, userAuth.UserId, providerID)
if err != nil {
util.WriteError(r.Context(), err, w)
return
}
models := make([]api.AgentNetworkDiscoveredModel, 0, len(result.Models))
for _, model := range result.Models {
models = append(models, api.AgentNetworkDiscoveredModel{
Id: model.ID,
Label: model.Label,
})
}
w.Header().Set("Cache-Control", "no-store")
util.WriteJSONObject(r.Context(), w, api.AgentNetworkModelDiscoveryResponse{
Models: models,
Source: api.AgentNetworkModelDiscoveryResponseSource(result.Source),
ProxyCluster: result.ProxyCluster,
RequestId: result.RequestID,
})
}
func (h *handler) createProvider(w http.ResponseWriter, r *http.Request) {
userAuth, err := nbcontext.GetUserAuthFromContext(r.Context())
if err != nil {
@@ -193,11 +229,12 @@ func (h *handler) deleteProvider(w http.ResponseWriter, r *http.Request) {
util.WriteJSONObject(r.Context(), w, util.EmptyObject{})
}
func validate(req *api.AgentNetworkProviderRequest, requireAPIKey bool) error {
func validate(req *api.AgentNetworkProviderRequest, creating bool) error {
if strings.TrimSpace(req.ProviderId) == "" {
return status.Errorf(status.InvalidArgument, "provider_id is required")
}
if !catalog.IsKnown(req.ProviderId) {
entry, ok := catalog.Lookup(req.ProviderId)
if !ok {
return status.Errorf(status.InvalidArgument, "provider_id %q is not a known catalog provider", req.ProviderId)
}
if strings.TrimSpace(req.Name) == "" {
@@ -210,8 +247,27 @@ func validate(req *api.AgentNetworkProviderRequest, requireAPIKey bool) error {
if err != nil || u.Host == "" || (u.Scheme != "http" && u.Scheme != "https") {
return status.Errorf(status.InvalidArgument, "upstream_url must be a full http(s) URL")
}
if requireAPIKey && (req.ApiKey == nil || strings.TrimSpace(*req.ApiKey) == "") {
return status.Errorf(status.InvalidArgument, "api_key is required")
return validateProviderAPIKey(entry, req.ApiKey, creating)
}
func validateProviderAPIKey(entry catalog.Provider, apiKey *string, creating bool) error {
apiKeyBlank := apiKey == nil || strings.TrimSpace(*apiKey) == ""
switch entry.EffectiveAuthMode() {
case catalog.AuthModeRequired:
// Updates may omit the key to preserve it, but an explicitly empty
// value cannot clear a credential required by the selected provider.
if (creating || apiKey != nil) && apiKeyBlank {
return status.Errorf(status.InvalidArgument, "api_key is required for provider_id %q", entry.ID)
}
case catalog.AuthModeOptional:
// Both omitted and explicitly empty values are valid. The manager
// distinguishes preserve from clear during update.
case catalog.AuthModeNone:
if !apiKeyBlank {
return status.Errorf(status.InvalidArgument, "api_key is not supported for provider_id %q", entry.ID)
}
default:
return status.Errorf(status.InvalidArgument, "provider_id %q has an invalid authentication mode", entry.ID)
}
return nil
}

View File

@@ -0,0 +1,47 @@
package handlers
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog"
)
func TestValidateProviderAPIKey(t *testing.T) {
key := "secret"
empty := ""
tests := []struct {
name string
mode catalog.AuthMode
apiKey *string
creating bool
wantErr string
}{
{name: "required create with key", mode: catalog.AuthModeRequired, apiKey: &key, creating: true},
{name: "required create omitted", mode: catalog.AuthModeRequired, creating: true, wantErr: "required"},
{name: "required update omitted preserves", mode: catalog.AuthModeRequired},
{name: "required update explicit empty", mode: catalog.AuthModeRequired, apiKey: &empty, wantErr: "required"},
{name: "optional create omitted", mode: catalog.AuthModeOptional, creating: true},
{name: "optional update explicit empty", mode: catalog.AuthModeOptional, apiKey: &empty},
{name: "optional with key", mode: catalog.AuthModeOptional, apiKey: &key, creating: true},
{name: "none omitted", mode: catalog.AuthModeNone, creating: true},
{name: "none explicit empty", mode: catalog.AuthModeNone, apiKey: &empty},
{name: "none with key", mode: catalog.AuthModeNone, apiKey: &key, wantErr: "not supported"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
entry := catalog.Provider{ID: "test-provider", AuthMode: tt.mode}
err := validateProviderAPIKey(entry, tt.apiKey, tt.creating)
if tt.wantErr == "" {
require.NoError(t, err)
return
}
require.Error(t, err)
assert.Contains(t, err.Error(), tt.wantErr)
})
}
}

View File

@@ -9,9 +9,12 @@ import (
"strings"
"sync"
"time"
"unicode"
"unicode/utf8"
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/labelgen"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
@@ -48,6 +51,7 @@ func ensureSessionKeys(p *types.Provider) error {
type Manager interface {
GetAllProviders(ctx context.Context, accountID, userID string) ([]*types.Provider, error)
GetProvider(ctx context.Context, accountID, userID, providerID string) (*types.Provider, error)
DiscoverProviderModels(ctx context.Context, accountID, userID, providerID string) (*types.ModelDiscoveryResult, error)
CreateProvider(ctx context.Context, userID string, provider *types.Provider, bootstrapCluster string) (*types.Provider, error)
UpdateProvider(ctx context.Context, userID string, provider *types.Provider) (*types.Provider, error)
DeleteProvider(ctx context.Context, accountID, userID, providerID string) error
@@ -128,8 +132,27 @@ type managerImpl struct {
// state; concurrent provider creates would otherwise race.
labelRngMu sync.Mutex
labelRng *rand.Rand
// discoveryAttempts provides lightweight per-provider admission control for
// the explicit endpoint probe. It prevents duplicate clicks from occupying
// multiple proxy control-stream requests at once and adds a short cooldown
// after each attempt.
discoveryMu sync.Mutex
discoveryAttempts map[string]modelDiscoveryAttempt
}
type modelDiscoveryAttempt struct {
inFlight bool
lastStarted time.Time
}
const (
modelDiscoveryTimeout = 10 * time.Second
modelDiscoveryCooldown = 2 * time.Second
maxDiscoveredModels = 500
maxDiscoveredModelLen = 512
)
// NewManager constructs the persistent Agent Network manager. The
// manager persists provider/policy/guardrail configuration and, on
// every mutation, reconciles the in-memory synthesised reverse-proxy
@@ -148,6 +171,7 @@ func NewManager(
proxyController: proxyController,
reconcileCache: make(map[string]map[string]*proto.ProxyMapping),
labelRng: rand.New(rand.NewSource(time.Now().UnixNano())),
discoveryAttempts: make(map[string]modelDiscoveryAttempt),
}
}
@@ -165,6 +189,202 @@ func (m *managerImpl) GetProvider(ctx context.Context, accountID, userID, provid
return m.store.GetAgentNetworkProviderByID(ctx, store.LockingStrengthNone, accountID, providerID)
}
// DiscoverProviderModels asks a capable proxy in the account's selected
// cluster to query the persisted provider endpoint. The browser supplies only
// the provider id: URL, TLS policy, and credential are loaded here so a caller
// cannot turn this operation into an arbitrary network probe.
func (m *managerImpl) DiscoverProviderModels(ctx context.Context, accountID, userID, providerID string) (*types.ModelDiscoveryResult, error) {
if err := m.requirePermission(ctx, accountID, userID, operations.Update); err != nil {
return nil, err
}
providerID = strings.TrimSpace(providerID)
if providerID == "" {
return nil, status.Errorf(status.InvalidArgument, "provider ID is required")
}
provider, err := m.store.GetAgentNetworkProviderByID(ctx, store.LockingStrengthNone, accountID, providerID)
if err != nil {
return nil, err
}
entry, ok := catalog.Lookup(provider.ProviderID)
if !ok {
return nil, status.Errorf(status.PreconditionFailed, "provider references an unknown catalog provider")
}
if entry.ModelDiscovery == nil {
return nil, status.Errorf(status.PreconditionFailed, "provider type %q does not support model discovery", provider.ProviderID)
}
if m.proxyController == nil {
return nil, status.Errorf(status.PreconditionFailed, "model discovery is unavailable")
}
settings, err := m.store.GetAgentNetworkSettings(ctx, store.LockingStrengthNone, accountID)
if err != nil {
var statusErr *status.Error
if errors.As(err, &statusErr) && statusErr.Type() == status.NotFound {
return nil, status.Errorf(status.PreconditionFailed, "configure an Agent Network proxy cluster before discovering models")
}
return nil, err
}
cluster := strings.TrimSpace(settings.Cluster)
if cluster == "" {
return nil, status.Errorf(status.PreconditionFailed, "configure an Agent Network proxy cluster before discovering models")
}
authHeaderName, authHeaderValue, gcpKey, err := providerAuthHeader(provider)
if err != nil {
return nil, status.Errorf(status.PreconditionFailed, "provider authentication is not configured correctly")
}
if gcpKey != "" {
return nil, status.Errorf(status.PreconditionFailed, "provider authentication is not supported for model discovery")
}
attemptKey := accountID + "\x00" + provider.ID
if !m.beginModelDiscovery(attemptKey, time.Now()) {
return nil, status.Errorf(status.TooManyRequests, "model discovery is already running or was requested too recently")
}
defer m.finishModelDiscovery(attemptKey)
probeCtx, cancel := context.WithTimeout(ctx, modelDiscoveryTimeout)
defer cancel()
probeResult, err := m.proxyController.DiscoverModels(probeCtx, accountID, cluster, &proto.ModelDiscoveryRequest{
UpstreamUrl: strings.TrimSpace(provider.UpstreamURL),
AuthHeaderName: authHeaderName,
AuthHeaderValue: authHeaderValue,
SkipTlsVerify: provider.SkipTLSVerification,
OllamaFallback: entry.ModelDiscovery.OllamaFallback,
})
if err != nil {
if errors.Is(err, context.DeadlineExceeded) || errors.Is(probeCtx.Err(), context.DeadlineExceeded) {
return nil, status.Errorf(status.PreconditionFailed, "model discovery timed out")
}
if _, ok := status.FromError(err); ok {
return nil, err
}
return nil, status.Errorf(status.PreconditionFailed, "model discovery could not be completed")
}
if probeResult == nil {
return nil, status.Errorf(status.PreconditionFailed, "proxy returned an empty model discovery response")
}
if strings.TrimSpace(probeResult.Error) != "" {
message := safeModelDiscoveryText(probeResult.Error, 256)
if message == "" {
message = "proxy reported a discovery failure"
}
requestID := safeModelDiscoveryText(probeResult.RequestId, 128)
if requestID != "" {
return nil, status.Errorf(status.PreconditionFailed, "model discovery failed: %s (request_id: %s)", message, requestID)
}
return nil, status.Errorf(status.PreconditionFailed, "model discovery failed: %s", message)
}
requestID := safeModelDiscoveryText(probeResult.RequestId, 128)
if requestID == "" {
return nil, status.Errorf(status.PreconditionFailed, "proxy returned an uncorrelated model discovery response")
}
source := strings.TrimSpace(probeResult.Source)
switch source {
case "openai_v1_models", "ollama_api_tags":
default:
return nil, status.Errorf(status.PreconditionFailed, "proxy returned an unsupported model discovery response")
}
return &types.ModelDiscoveryResult{
RequestID: requestID,
Source: source,
ProxyCluster: cluster,
Models: normalizeDiscoveredModels(probeResult.Models),
}, nil
}
func (m *managerImpl) beginModelDiscovery(key string, now time.Time) bool {
m.discoveryMu.Lock()
defer m.discoveryMu.Unlock()
if m.discoveryAttempts == nil {
m.discoveryAttempts = make(map[string]modelDiscoveryAttempt)
}
attempt := m.discoveryAttempts[key]
if attempt.inFlight || (!attempt.lastStarted.IsZero() && now.Sub(attempt.lastStarted) < modelDiscoveryCooldown) {
return false
}
attempt.inFlight = true
attempt.lastStarted = now
m.discoveryAttempts[key] = attempt
return true
}
func (m *managerImpl) finishModelDiscovery(key string) {
m.discoveryMu.Lock()
attempt, ok := m.discoveryAttempts[key]
if !ok {
m.discoveryMu.Unlock()
return
}
attempt.inFlight = false
m.discoveryAttempts[key] = attempt
m.discoveryMu.Unlock()
remaining := time.Until(attempt.lastStarted.Add(modelDiscoveryCooldown))
if remaining <= 0 {
m.expireModelDiscoveryAttempt(key, attempt.lastStarted)
return
}
time.AfterFunc(remaining, func() {
m.expireModelDiscoveryAttempt(key, attempt.lastStarted)
})
}
func (m *managerImpl) expireModelDiscoveryAttempt(key string, lastStarted time.Time) {
m.discoveryMu.Lock()
defer m.discoveryMu.Unlock()
attempt, ok := m.discoveryAttempts[key]
if ok && !attempt.inFlight && attempt.lastStarted.Equal(lastStarted) {
delete(m.discoveryAttempts, key)
}
}
func normalizeDiscoveredModels(models []*proto.ModelDiscoveryModel) []types.DiscoveredModel {
out := make([]types.DiscoveredModel, 0, min(len(models), maxDiscoveredModels))
seen := make(map[string]struct{}, min(len(models), maxDiscoveredModels))
for _, model := range models {
if model == nil {
continue
}
id := safeModelDiscoveryText(model.Id, maxDiscoveredModelLen)
if id == "" || len(id) > maxDiscoveredModelLen {
continue
}
if _, ok := seen[id]; ok {
continue
}
label := safeModelDiscoveryText(model.Label, maxDiscoveredModelLen)
if label == "" || len(label) > maxDiscoveredModelLen {
label = id
}
seen[id] = struct{}{}
out = append(out, types.DiscoveredModel{ID: id, Label: label})
if len(out) == maxDiscoveredModels {
break
}
}
slices.SortFunc(out, func(a, b types.DiscoveredModel) int {
return strings.Compare(a.ID, b.ID)
})
return out
}
func safeModelDiscoveryText(value string, maxLen int) string {
value = strings.TrimSpace(value)
if value == "" || len(value) > maxLen || !utf8.ValidString(value) {
return ""
}
for _, r := range value {
if unicode.IsControl(r) {
return ""
}
}
return value
}
// CreateProvider persists a new provider for the account. bootstrapCluster
// is used only when the per-account agent-network Settings row hasn't
// been created yet; otherwise it is ignored (the cluster is pinned on
@@ -174,11 +394,8 @@ func (m *managerImpl) CreateProvider(ctx context.Context, userID string, provide
return nil, err
}
// An empty api_key would silently produce a synthesised service
// that 401s on every upstream request. Surface the misconfiguration
// at create time instead.
if strings.TrimSpace(provider.APIKey) == "" {
return nil, status.Errorf(status.InvalidArgument, "api_key is required when creating an agent network provider")
if err := prepareProviderAPIKey(provider, nil); err != nil {
return nil, err
}
if provider.ID == "" {
@@ -222,13 +439,8 @@ func (m *managerImpl) UpdateProvider(ctx context.Context, userID string, provide
return nil, fmt.Errorf("failed to get agent network provider: %w", err)
}
// Preserve the API key if the caller didn't rotate it. A
// whitespace-only value is treated as "not rotated" rather than a
// real key, but it must not silently overwrite a valid stored key.
if provider.APIKey == "" {
provider.APIKey = existing.APIKey
} else if strings.TrimSpace(provider.APIKey) == "" {
return nil, status.Errorf(status.InvalidArgument, "api_key must be non-blank when rotating an agent network provider")
if err := prepareProviderAPIKey(provider, existing); err != nil {
return nil, err
}
// Always preserve the session keypair across updates so existing
// session cookies stay valid. The keys are server-managed and
@@ -251,6 +463,52 @@ func (m *managerImpl) UpdateProvider(ctx context.Context, userID string, provide
return provider, nil
}
// prepareProviderAPIKey applies catalog-owned authentication semantics before
// a provider is persisted. existing is nil on create. On update, omission
// preserves a key only when the provider type is unchanged; switching types
// never carries an old provider's secret into the new upstream.
func prepareProviderAPIKey(provider, existing *types.Provider) error {
entry, ok := catalog.Lookup(provider.ProviderID)
if !ok {
return status.Errorf(status.InvalidArgument, "provider_id %q is not a known catalog provider", provider.ProviderID)
}
return prepareProviderAPIKeyForEntry(provider, existing, entry)
}
func prepareProviderAPIKeyForEntry(provider, existing *types.Provider, entry catalog.Provider) error {
keyProvided := provider.APIKeyProvided || provider.APIKey != ""
providerChanged := existing != nil && provider.ProviderID != existing.ProviderID
if existing != nil && !keyProvided && !providerChanged && entry.EffectiveAuthMode() != catalog.AuthModeNone {
provider.APIKey = existing.APIKey
}
// Normalize after a preserved value is restored so a legacy
// whitespace-only credential cannot pass manager validation and then fail
// later during synthesis.
if strings.TrimSpace(provider.APIKey) == "" {
provider.APIKey = ""
}
switch entry.EffectiveAuthMode() {
case catalog.AuthModeRequired:
if provider.APIKey == "" {
return status.Errorf(status.InvalidArgument, "api_key is required for provider_id %q", provider.ProviderID)
}
case catalog.AuthModeOptional:
// Empty is a valid credential state.
case catalog.AuthModeNone:
if keyProvided && provider.APIKey != "" {
return status.Errorf(status.InvalidArgument, "api_key is not supported for provider_id %q", provider.ProviderID)
}
// Clear any stale credential already stored for a provider whose
// catalog semantics changed to no authentication.
provider.APIKey = ""
default:
return status.Errorf(status.InvalidArgument, "provider_id %q has an invalid authentication mode", provider.ProviderID)
}
return nil
}
func (m *managerImpl) DeleteProvider(ctx context.Context, accountID, userID, providerID string) error {
if err := m.requirePermission(ctx, accountID, userID, operations.Delete); err != nil {
return err
@@ -808,6 +1066,10 @@ func (*mockManager) GetProvider(_ context.Context, _, _, _ string) (*types.Provi
return &types.Provider{}, nil
}
func (*mockManager) DiscoverProviderModels(_ context.Context, _, _, _ string) (*types.ModelDiscoveryResult, error) {
return nil, status.Errorf(status.PreconditionFailed, "model discovery is unavailable")
}
func (*mockManager) CreateProvider(_ context.Context, _ string, p *types.Provider, _ string) (*types.Provider, error) {
return p, nil
}

View File

@@ -0,0 +1,248 @@
package agentnetwork
import (
"context"
"strings"
"testing"
"time"
"github.com/golang/mock/gomock"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
"github.com/netbirdio/netbird/management/server/permissions"
"github.com/netbirdio/netbird/management/server/permissions/modules"
"github.com/netbirdio/netbird/management/server/permissions/operations"
"github.com/netbirdio/netbird/management/server/store"
"github.com/netbirdio/netbird/shared/management/proto"
"github.com/netbirdio/netbird/shared/management/status"
)
func newModelDiscoveryManager(t *testing.T) (*managerImpl, *store.MockStore, *permissions.MockManager, *proxy.MockController) {
t.Helper()
ctrl := gomock.NewController(t)
mockStore := store.NewMockStore(ctrl)
mockPermissions := permissions.NewMockManager(ctrl)
mockProxy := proxy.NewMockController(ctrl)
return &managerImpl{
store: mockStore,
permissionsManager: mockPermissions,
proxyController: mockProxy,
discoveryAttempts: make(map[string]modelDiscoveryAttempt),
}, mockStore, mockPermissions, mockProxy
}
func allowModelDiscovery(mockPermissions *permissions.MockManager) {
mockPermissions.EXPECT().
ValidateUserPermissions(gomock.Any(), "account-1", "user-1", modules.AgentNetwork, operations.Update).
Return(true, context.Background(), nil)
}
func TestDiscoverProviderModelsUsesPersistedProviderAndCluster(t *testing.T) {
manager, mockStore, mockPermissions, mockProxy := newModelDiscoveryManager(t)
allowModelDiscovery(mockPermissions)
provider := &types.Provider{
ID: "provider-1",
AccountID: "account-1",
ProviderID: "ollama",
UpstreamURL: "http://ollama.internal:11434/base",
APIKey: "stored-key",
SkipTLSVerification: true,
}
mockStore.EXPECT().
GetAgentNetworkProviderByID(gomock.Any(), store.LockingStrengthNone, "account-1", "provider-1").
Return(provider, nil)
mockStore.EXPECT().
GetAgentNetworkSettings(gomock.Any(), store.LockingStrengthNone, "account-1").
Return(&types.Settings{AccountID: "account-1", Cluster: "private.example.com"}, nil)
mockProxy.EXPECT().
DiscoverModels(gomock.Any(), "account-1", "private.example.com", gomock.Any()).
DoAndReturn(func(_ context.Context, _, _ string, request *proto.ModelDiscoveryRequest) (*proto.ModelDiscoveryResult, error) {
assert.Empty(t, request.RequestId, "the control-plane server owns request IDs")
assert.Equal(t, provider.UpstreamURL, request.UpstreamUrl)
assert.Equal(t, "Authorization", request.AuthHeaderName)
assert.Equal(t, "Bearer stored-key", request.AuthHeaderValue)
assert.True(t, request.SkipTlsVerify)
assert.True(t, request.OllamaFallback)
return &proto.ModelDiscoveryResult{
RequestId: "probe-123",
Source: "openai_v1_models",
Models: []*proto.ModelDiscoveryModel{
{Id: "zeta", Label: ""},
{Id: "alpha", Label: "Alpha"},
{Id: "zeta", Label: "duplicate"},
{Id: "bad\x00model", Label: "bad"},
{Id: strings.Repeat("x", maxDiscoveredModelLen+1), Label: "too long"},
},
}, nil
})
result, err := manager.DiscoverProviderModels(context.Background(), "account-1", "user-1", "provider-1")
require.NoError(t, err)
assert.Equal(t, "probe-123", result.RequestID)
assert.Equal(t, "openai_v1_models", result.Source)
assert.Equal(t, "private.example.com", result.ProxyCluster)
assert.Equal(t, []types.DiscoveredModel{
{ID: "alpha", Label: "Alpha"},
{ID: "zeta", Label: "zeta"},
}, result.Models)
}
func TestDiscoverProviderModelsSupportsOllamaCloud(t *testing.T) {
manager, mockStore, mockPermissions, mockProxy := newModelDiscoveryManager(t)
allowModelDiscovery(mockPermissions)
provider := &types.Provider{
ID: "provider-cloud",
AccountID: "account-1",
ProviderID: "ollama_cloud",
UpstreamURL: "https://ollama.com",
APIKey: "ollama-cloud-key",
}
mockStore.EXPECT().
GetAgentNetworkProviderByID(gomock.Any(), store.LockingStrengthNone, "account-1", provider.ID).
Return(provider, nil)
mockStore.EXPECT().
GetAgentNetworkSettings(gomock.Any(), store.LockingStrengthNone, "account-1").
Return(&types.Settings{AccountID: "account-1", Cluster: "cloud-proxies.example.com"}, nil)
mockProxy.EXPECT().
DiscoverModels(gomock.Any(), "account-1", "cloud-proxies.example.com", gomock.Any()).
DoAndReturn(func(_ context.Context, _, _ string, request *proto.ModelDiscoveryRequest) (*proto.ModelDiscoveryResult, error) {
assert.Equal(t, "https://ollama.com", request.UpstreamUrl)
assert.Equal(t, "Authorization", request.AuthHeaderName)
assert.Equal(t, "Bearer ollama-cloud-key", request.AuthHeaderValue)
assert.False(t, request.SkipTlsVerify)
assert.True(t, request.OllamaFallback)
return &proto.ModelDiscoveryResult{
RequestId: "probe-cloud",
Source: "openai_v1_models",
Models: []*proto.ModelDiscoveryModel{
{Id: "gpt-oss:120b", Label: "gpt-oss:120b"},
},
}, nil
})
result, err := manager.DiscoverProviderModels(context.Background(), "account-1", "user-1", provider.ID)
require.NoError(t, err)
assert.Equal(t, "probe-cloud", result.RequestID)
assert.Equal(t, "cloud-proxies.example.com", result.ProxyCluster)
assert.Equal(t, []types.DiscoveredModel{
{ID: "gpt-oss:120b", Label: "gpt-oss:120b"},
}, result.Models)
}
func TestDiscoverProviderModelsRejectsOllamaCloudWithoutKey(t *testing.T) {
manager, mockStore, mockPermissions, _ := newModelDiscoveryManager(t)
allowModelDiscovery(mockPermissions)
provider := &types.Provider{
ID: "provider-cloud",
AccountID: "account-1",
ProviderID: "ollama_cloud",
UpstreamURL: "https://ollama.com",
}
mockStore.EXPECT().
GetAgentNetworkProviderByID(gomock.Any(), store.LockingStrengthNone, "account-1", provider.ID).
Return(provider, nil)
mockStore.EXPECT().
GetAgentNetworkSettings(gomock.Any(), store.LockingStrengthNone, "account-1").
Return(&types.Settings{AccountID: "account-1", Cluster: "cloud-proxies.example.com"}, nil)
_, err := manager.DiscoverProviderModels(context.Background(), "account-1", "user-1", provider.ID)
require.Error(t, err)
statusErr, ok := status.FromError(err)
require.True(t, ok)
assert.Equal(t, status.PreconditionFailed, statusErr.Type())
assert.Contains(t, err.Error(), "authentication is not configured correctly")
}
func TestDiscoverProviderModelsRejectsUnsupportedProvider(t *testing.T) {
manager, mockStore, mockPermissions, _ := newModelDiscoveryManager(t)
allowModelDiscovery(mockPermissions)
mockStore.EXPECT().
GetAgentNetworkProviderByID(gomock.Any(), store.LockingStrengthNone, "account-1", "provider-1").
Return(&types.Provider{
ID: "provider-1",
AccountID: "account-1",
ProviderID: "openai_api",
}, nil)
_, err := manager.DiscoverProviderModels(context.Background(), "account-1", "user-1", "provider-1")
require.Error(t, err)
statusErr, ok := status.FromError(err)
require.True(t, ok)
assert.Equal(t, status.PreconditionFailed, statusErr.Type())
assert.Contains(t, err.Error(), "does not support")
}
func TestDiscoverProviderModelsPreservesCorrelatedProxyError(t *testing.T) {
manager, mockStore, mockPermissions, mockProxy := newModelDiscoveryManager(t)
allowModelDiscovery(mockPermissions)
mockStore.EXPECT().
GetAgentNetworkProviderByID(gomock.Any(), store.LockingStrengthNone, "account-1", "provider-1").
Return(&types.Provider{
ID: "provider-1",
AccountID: "account-1",
ProviderID: "ollama",
UpstreamURL: "http://ollama.internal:11434",
}, nil)
mockStore.EXPECT().
GetAgentNetworkSettings(gomock.Any(), store.LockingStrengthNone, "account-1").
Return(&types.Settings{AccountID: "account-1", Cluster: "private.example.com"}, nil)
mockProxy.EXPECT().
DiscoverModels(gomock.Any(), "account-1", "private.example.com", gomock.Any()).
Return(&proto.ModelDiscoveryResult{
RequestId: "probe-failed",
Error: "upstream returned HTTP 401",
}, nil)
_, err := manager.DiscoverProviderModels(context.Background(), "account-1", "user-1", "provider-1")
require.Error(t, err)
assert.Contains(t, err.Error(), "upstream returned HTTP 401")
assert.Contains(t, err.Error(), "probe-failed")
}
func TestModelDiscoveryAdmissionControl(t *testing.T) {
manager := &managerImpl{}
start := time.Now()
require.True(t, manager.beginModelDiscovery("account/provider", start))
assert.False(t, manager.beginModelDiscovery("account/provider", start.Add(time.Second)))
manager.finishModelDiscovery("account/provider")
assert.False(t, manager.beginModelDiscovery("account/provider", start.Add(time.Second)))
assert.True(t, manager.beginModelDiscovery("account/provider", start.Add(modelDiscoveryCooldown)))
}
func TestModelDiscoveryAdmissionControlExpiresFinishedAttempts(t *testing.T) {
manager := &managerImpl{}
start := time.Now()
require.True(t, manager.beginModelDiscovery("account/provider", start))
manager.finishModelDiscovery("account/provider")
manager.expireModelDiscoveryAttempt("account/provider", start)
manager.discoveryMu.Lock()
_, retained := manager.discoveryAttempts["account/provider"]
manager.discoveryMu.Unlock()
assert.False(t, retained)
}
func TestModelDiscoveryAdmissionControlKeepsReplacementAttempt(t *testing.T) {
manager := &managerImpl{}
start := time.Now()
require.True(t, manager.beginModelDiscovery("account/provider", start))
manager.finishModelDiscovery("account/provider")
require.True(t, manager.beginModelDiscovery("account/provider", start.Add(modelDiscoveryCooldown)))
manager.expireModelDiscoveryAttempt("account/provider", start)
manager.discoveryMu.Lock()
attempt, retained := manager.discoveryAttempts["account/provider"]
manager.discoveryMu.Unlock()
require.True(t, retained)
assert.True(t, attempt.inFlight)
assert.Equal(t, start.Add(modelDiscoveryCooldown), attempt.lastStarted)
}

View File

@@ -0,0 +1,124 @@
package agentnetwork
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
)
func TestPrepareProviderAPIKey(t *testing.T) {
t.Run("required create rejects an empty key", func(t *testing.T) {
provider := &types.Provider{ProviderID: "openai_api"}
err := prepareProviderAPIKey(provider, nil)
require.Error(t, err)
assert.Contains(t, err.Error(), "api_key is required")
})
t.Run("Ollama Cloud requires a key", func(t *testing.T) {
provider := &types.Provider{ProviderID: "ollama_cloud"}
err := prepareProviderAPIKey(provider, nil)
require.Error(t, err)
assert.Contains(t, err.Error(), "api_key is required")
})
t.Run("Ollama Cloud accepts a key", func(t *testing.T) {
provider := &types.Provider{
ProviderID: "ollama_cloud",
APIKey: "ollama-cloud-key",
APIKeyProvided: true,
}
require.NoError(t, prepareProviderAPIKey(provider, nil))
assert.Equal(t, "ollama-cloud-key", provider.APIKey)
})
t.Run("optional create accepts an empty key", func(t *testing.T) {
provider := &types.Provider{ProviderID: "ollama"}
require.NoError(t, prepareProviderAPIKey(provider, nil))
assert.Empty(t, provider.APIKey)
})
t.Run("optional update omission preserves the existing key", func(t *testing.T) {
existing := &types.Provider{ProviderID: "ollama", APIKey: "existing"}
provider := &types.Provider{ProviderID: "ollama"}
require.NoError(t, prepareProviderAPIKey(provider, existing))
assert.Equal(t, "existing", provider.APIKey)
})
t.Run("optional update explicit empty clears the existing key", func(t *testing.T) {
existing := &types.Provider{ProviderID: "ollama", APIKey: "existing"}
provider := &types.Provider{ProviderID: "ollama", APIKeyProvided: true}
require.NoError(t, prepareProviderAPIKey(provider, existing))
assert.Empty(t, provider.APIKey)
})
t.Run("changing to optional without a key does not carry the old secret", func(t *testing.T) {
existing := &types.Provider{ProviderID: "openai_api", APIKey: "sk-old-provider"}
provider := &types.Provider{ProviderID: "ollama"}
require.NoError(t, prepareProviderAPIKey(provider, existing))
assert.Empty(t, provider.APIKey)
})
t.Run("changing to required without a key fails", func(t *testing.T) {
existing := &types.Provider{ProviderID: "ollama", APIKey: "old-optional-key"}
provider := &types.Provider{ProviderID: "openai_api"}
err := prepareProviderAPIKey(provider, existing)
require.Error(t, err)
assert.Contains(t, err.Error(), "api_key is required")
assert.Empty(t, provider.APIKey)
})
t.Run("required update omission preserves the existing key", func(t *testing.T) {
existing := &types.Provider{ProviderID: "openai_api", APIKey: "sk-existing"}
provider := &types.Provider{ProviderID: "openai_api"}
require.NoError(t, prepareProviderAPIKey(provider, existing))
assert.Equal(t, "sk-existing", provider.APIKey)
})
t.Run("required update rejects a preserved whitespace-only key", func(t *testing.T) {
existing := &types.Provider{ProviderID: "openai_api", APIKey: " "}
provider := &types.Provider{ProviderID: "openai_api"}
err := prepareProviderAPIKey(provider, existing)
require.Error(t, err)
assert.Contains(t, err.Error(), "api_key is required")
assert.Empty(t, provider.APIKey)
})
t.Run("required update explicit empty is rejected", func(t *testing.T) {
existing := &types.Provider{ProviderID: "openai_api", APIKey: "sk-existing"}
provider := &types.Provider{ProviderID: "openai_api", APIKeyProvided: true}
err := prepareProviderAPIKey(provider, existing)
require.Error(t, err)
assert.Contains(t, err.Error(), "api_key is required")
})
t.Run("unknown catalog provider is rejected", func(t *testing.T) {
provider := &types.Provider{ProviderID: "unknown", APIKey: "secret"}
err := prepareProviderAPIKey(provider, nil)
require.Error(t, err)
assert.Contains(t, err.Error(), "not a known catalog provider")
})
t.Run("none mode clears a stale key", func(t *testing.T) {
entry := catalog.Provider{ID: "authless", AuthMode: catalog.AuthModeNone}
existing := &types.Provider{ProviderID: "authless", APIKey: "stale-secret"}
provider := &types.Provider{ProviderID: "authless"}
require.NoError(t, prepareProviderAPIKeyForEntry(provider, existing, entry))
assert.Empty(t, provider.APIKey)
})
t.Run("none mode rejects a supplied key", func(t *testing.T) {
entry := catalog.Provider{ID: "authless", AuthMode: catalog.AuthModeNone}
provider := &types.Provider{
ProviderID: "authless",
APIKey: "unexpected-secret",
APIKeyProvided: true,
}
err := prepareProviderAPIKeyForEntry(provider, nil, entry)
require.Error(t, err)
assert.Contains(t, err.Error(), "not supported")
})
}

View File

@@ -906,23 +906,33 @@ const (
noopUpstreamPort = uint16(443)
)
// providerAuthHeader builds the upstream auth header pair for a
// provider from its catalog entry. The catalog declares which header
// name and template a provider's API expects; the synthesiser
// substitutes the provider's decrypted API key into the template and
// returns the (name, value) pair the router middleware injects after
// stripping the inbound vendor auth headers.
// providerAuthHeader builds the upstream auth header pair for a provider from
// its catalog entry. Optional providers with no key and providers using no
// authentication return an empty pair. Otherwise, the synthesiser substitutes
// the decrypted API key into the catalog template and returns the pair the
// router injects after stripping inbound vendor credentials.
func providerAuthHeader(p *types.Provider) (name, value, gcpSAKeyB64 string, err error) {
entry, ok := catalog.Lookup(p.ProviderID)
if !ok {
return "", "", "", fmt.Errorf("provider %s references unknown catalog id %q", p.ID, p.ProviderID)
}
switch entry.EffectiveAuthMode() {
case catalog.AuthModeNone:
return "", "", "", nil
case catalog.AuthModeOptional:
if strings.TrimSpace(p.APIKey) == "" {
return "", "", "", nil
}
case catalog.AuthModeRequired:
if strings.TrimSpace(p.APIKey) == "" {
return "", "", "", fmt.Errorf("provider %s has no api key", p.ID)
}
default:
return "", "", "", fmt.Errorf("catalog entry %q has invalid authentication mode %q", p.ProviderID, entry.AuthMode)
}
if entry.AuthHeaderName == "" || entry.AuthHeaderTemplate == "" {
return "", "", "", fmt.Errorf("catalog entry %q has no auth header configured", p.ProviderID)
}
if p.APIKey == "" {
return "", "", "", fmt.Errorf("provider %s has no api key", p.ID)
}
// A "keyfile::<base64 json>" api_key is a GCP service-account key, not a
// static bearer. The proxy mints + refreshes a short-lived OAuth token from
// it at request time, so carry the key material on the route and emit no

View File

@@ -1219,3 +1219,102 @@ func TestSynthesizeServices_EmptyAPIKey_FailsClosed(t *testing.T) {
require.Error(t, err, "synthesis must refuse a provider with no api key")
assert.Contains(t, err.Error(), "no api key", "error must surface the missing credential")
}
func TestSynthesizeServices_OllamaOptionalAPIKey(t *testing.T) {
tests := []struct {
name string
apiKey string
wantHeaderName string
wantHeaderValue string
}{
{
name: "no key emits no replacement auth header",
},
{
name: "configured key emits bearer auth",
apiKey: "protected-endpoint-token",
wantHeaderName: "Authorization",
wantHeaderValue: "Bearer protected-endpoint-token",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ctx := context.Background()
ctrl := gomock.NewController(t)
defer ctrl.Finish()
mockStore := store.NewMockStore(ctrl)
provider := newSynthTestProvider()
provider.ProviderID = "ollama"
provider.Name = "Ollama"
provider.UpstreamURL = "http://ollama.internal:11434"
provider.APIKey = tt.apiKey
provider.Models = []types.ProviderModel{}
policy := newSynthTestPolicy(provider.ID, "grp-eng", "")
expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(),
[]*types.Provider{provider},
[]*types.Policy{policy},
[]*types.Guardrail{})
services, err := SynthesizeServices(ctx, mockStore, testAccountID)
require.NoError(t, err)
require.Len(t, services, 1)
var routerCfg routerConfig
for _, middleware := range services[0].Targets[0].Options.Middlewares {
if middleware.ID == middlewareIDLLMRouter {
require.NoError(t, json.Unmarshal(middleware.ConfigJSON, &routerCfg))
break
}
}
require.Len(t, routerCfg.Providers, 1)
assert.Equal(t, tt.wantHeaderName, routerCfg.Providers[0].AuthHeaderName)
assert.Equal(t, tt.wantHeaderValue, routerCfg.Providers[0].AuthHeaderValue)
})
}
}
func TestSynthesizeServices_OllamaCloudRoute(t *testing.T) {
ctx := context.Background()
ctrl := gomock.NewController(t)
defer ctrl.Finish()
mockStore := store.NewMockStore(ctrl)
provider := newSynthTestProvider()
provider.ProviderID = "ollama_cloud"
provider.Name = "Ollama Cloud"
provider.UpstreamURL = "https://ollama.com"
provider.APIKey = "ollama-cloud-key"
provider.Models = []types.ProviderModel{{ID: "gpt-oss:120b"}}
policy := newSynthTestPolicy(provider.ID, "grp-eng", "")
expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(),
[]*types.Provider{provider},
[]*types.Policy{policy},
[]*types.Guardrail{})
services, err := SynthesizeServices(ctx, mockStore, testAccountID)
require.NoError(t, err)
require.Len(t, services, 1)
var routerCfg routerConfig
for _, middleware := range services[0].Targets[0].Options.Middlewares {
if middleware.ID == middlewareIDLLMRouter {
require.NoError(t, json.Unmarshal(middleware.ConfigJSON, &routerCfg))
break
}
}
require.Len(t, routerCfg.Providers, 1)
route := routerCfg.Providers[0]
assert.Empty(t, route.Vendor, "Ollama Cloud preserves the untagged Ollama/vLLM/custom routing behavior")
assert.Equal(t, []string{"gpt-oss:120b"}, route.Models)
assert.Equal(t, "https", route.UpstreamScheme)
assert.Equal(t, "ollama.com", route.UpstreamHost)
assert.Empty(t, route.UpstreamPath)
assert.Equal(t, "Authorization", route.AuthHeaderName)
assert.Equal(t, "Bearer ollama-cloud-key", route.AuthHeaderValue)
assert.False(t, route.SkipTLSVerify)
}

View File

@@ -0,0 +1,19 @@
package types
// DiscoveredModel is a normalized model identifier returned by a provider's
// persisted upstream endpoint. Discovery does not persist models; the caller
// must explicitly update the provider to opt into the returned list.
type DiscoveredModel struct {
ID string
Label string
}
// ModelDiscoveryResult describes a successful proxy-executed discovery probe.
// RequestID correlates the management request with the proxy control message,
// and ProxyCluster identifies the network vantage point that ran it.
type ModelDiscoveryResult struct {
RequestID string
Source string
ProxyCluster string
Models []DiscoveredModel
}

View File

@@ -33,6 +33,10 @@ type Provider struct {
// the operator selected.
UpstreamURL string `gorm:"column:upstream_url"`
APIKey string `gorm:"column:api_key"`
// APIKeyProvided records whether api_key was present in the request. It is
// transient and lets updates distinguish omission (preserve) from an
// explicit empty value (clear an optional credential).
APIKeyProvided bool `gorm:"-" json:"-"`
// ExtraValues holds operator-typed values for catalog-declared
// ExtraHeaders (see catalog.Provider.ExtraHeaders). Keyed by
// header name (e.g. "x-portkey-config"); a non-empty value is
@@ -97,15 +101,16 @@ func NewProvider(accountID string) *Provider {
}
}
// FromAPIRequest applies the request payload onto the receiver. The api_key
// is only overwritten when the caller provided one — empty/nil leaves the
// existing key intact, so updates can omit it.
// FromAPIRequest applies the request payload onto the receiver. APIKeyProvided
// preserves the distinction between an omitted api_key and an explicit empty
// value; the manager applies the catalog-specific preserve/clear semantics.
func (p *Provider) FromAPIRequest(req *api.AgentNetworkProviderRequest) {
p.ProviderID = req.ProviderId
p.Name = req.Name
p.UpstreamURL = req.UpstreamUrl
if req.ApiKey != nil && strings.TrimSpace(*req.ApiKey) != "" {
p.APIKey = *req.ApiKey
p.APIKeyProvided = req.ApiKey != nil
if req.ApiKey != nil {
p.APIKey = strings.TrimSpace(*req.ApiKey)
}
if req.ExtraValues != nil {
// Replace the whole map (rather than merge) so unsetting a
@@ -178,6 +183,7 @@ func (p *Provider) ToAPIResponse() *api.AgentNetworkProvider {
UpstreamUrl: p.UpstreamURL,
Models: models,
Enabled: p.Enabled,
HasApiKey: strings.TrimSpace(p.APIKey) != "",
SkipTlsVerification: p.SkipTLSVerification,
MetadataDisabled: p.MetadataDisabled,
CreatedAt: &created,

View File

@@ -77,3 +77,45 @@ func TestProvider_MetadataDisabled_RoundTrip(t *testing.T) {
assert.False(t, p.MetadataDisabled, "explicit false must clear metadata_disabled")
assert.False(t, p.ToAPIResponse().MetadataDisabled, "response must reflect the cleared value")
}
func TestProvider_APIKeyPresenceAndResponse(t *testing.T) {
base := func() *api.AgentNetworkProviderRequest {
return &api.AgentNetworkProviderRequest{
ProviderId: "ollama",
Name: "Ollama",
UpstreamUrl: "http://ollama.internal:11434",
}
}
p := NewProvider("acc-1")
p.FromAPIRequest(base())
assert.False(t, p.APIKeyProvided, "omission must remain distinguishable from an explicit clear")
assert.False(t, p.ToAPIResponse().HasApiKey)
key := "protected-endpoint-token"
withKey := base()
withKey.ApiKey = &key
p.FromAPIRequest(withKey)
assert.True(t, p.APIKeyProvided)
assert.Equal(t, key, p.APIKey)
assert.True(t, p.ToAPIResponse().HasApiKey)
paddedKey := " \tprotected-endpoint-token\n"
withPaddedKey := base()
withPaddedKey.ApiKey = &paddedKey
p.FromAPIRequest(withPaddedKey)
assert.True(t, p.APIKeyProvided)
assert.Equal(t, key, p.APIKey, "non-blank API keys must be normalized before storage")
assert.True(t, p.ToAPIResponse().HasApiKey, "the response must reflect the normalized stored key")
empty := ""
clearKey := base()
clearKey.ApiKey = &empty
p.FromAPIRequest(clearKey)
assert.True(t, p.APIKeyProvided, "explicit empty must be carried to the manager as a clear")
assert.Empty(t, p.APIKey)
assert.False(t, p.ToAPIResponse().HasApiKey)
p.FromAPIRequest(base())
assert.False(t, p.APIKeyProvided, "a later omitted value must reset transient request presence")
}

View File

@@ -4,11 +4,16 @@ package proxy
import (
"context"
"errors"
"time"
"github.com/netbirdio/netbird/shared/management/proto"
)
// ErrModelDiscoveryUnavailable is returned when no connected proxy in the
// requested cluster can execute a model-discovery request.
var ErrModelDiscoveryUnavailable = errors.New("model discovery proxy unavailable")
// Manager defines the interface for proxy operations
type Manager interface {
Connect(ctx context.Context, proxyID, sessionID, clusterAddress, ipAddress string, accountID *string, capabilities *Capabilities) (*Proxy, error)
@@ -38,6 +43,7 @@ type OIDCValidationConfig struct {
// Controller is responsible for managing proxy clusters and routing service updates.
type Controller interface {
SendServiceUpdateToCluster(ctx context.Context, accountID string, update *proto.ProxyMapping, clusterAddr string)
DiscoverModels(ctx context.Context, accountID, clusterAddr string, req *proto.ModelDiscoveryRequest) (*proto.ModelDiscoveryResult, error)
GetOIDCValidationConfig() OIDCValidationConfig
RegisterProxyToCluster(ctx context.Context, clusterAddr, proxyID string) error
UnregisterProxyFromCluster(ctx context.Context, clusterAddr, proxyID string) error

View File

@@ -39,6 +39,12 @@ func (c *GRPCController) SendServiceUpdateToCluster(ctx context.Context, account
c.metrics.IncrementServiceUpdateSendCount(clusterAddr)
}
// DiscoverModels executes a correlated model-discovery request on one capable
// proxy connected to the requested cluster.
func (c *GRPCController) DiscoverModels(ctx context.Context, accountID, clusterAddr string, req *proto.ModelDiscoveryRequest) (*proto.ModelDiscoveryResult, error) {
return c.proxyGRPCServer.DiscoverModels(ctx, accountID, clusterAddr, req)
}
// GetOIDCValidationConfig returns the OIDC validation configuration from the gRPC server.
func (c *GRPCController) GetOIDCValidationConfig() proxy.OIDCValidationConfig {
return c.proxyGRPCServer.GetOIDCValidationConfig()
@@ -50,10 +56,12 @@ func (c *GRPCController) RegisterProxyToCluster(ctx context.Context, clusterAddr
return nil
}
proxySet, _ := c.clusterProxies.LoadOrStore(clusterAddr, &sync.Map{})
proxySet.(*sync.Map).Store(proxyID, struct{}{})
_, alreadyRegistered := proxySet.(*sync.Map).LoadOrStore(proxyID, struct{}{})
log.WithContext(ctx).Debugf("Registered proxy %s to cluster %s", proxyID, clusterAddr)
c.metrics.IncrementProxyConnectionCount(clusterAddr)
if !alreadyRegistered {
c.metrics.IncrementProxyConnectionCount(clusterAddr)
}
return nil
}
@@ -64,10 +72,10 @@ func (c *GRPCController) UnregisterProxyFromCluster(ctx context.Context, cluster
return nil
}
if proxySet, ok := c.clusterProxies.Load(clusterAddr); ok {
proxySet.(*sync.Map).Delete(proxyID)
log.WithContext(ctx).Debugf("Unregistered proxy %s from cluster %s", proxyID, clusterAddr)
c.metrics.DecrementProxyConnectionCount(clusterAddr)
if _, registered := proxySet.(*sync.Map).LoadAndDelete(proxyID); registered {
log.WithContext(ctx).Debugf("Unregistered proxy %s from cluster %s", proxyID, clusterAddr)
c.metrics.DecrementProxyConnectionCount(clusterAddr)
}
}
return nil
}

View File

@@ -259,6 +259,21 @@ func (m *MockController) EXPECT() *MockControllerMockRecorder {
return m.recorder
}
// DiscoverModels mocks base method.
func (m *MockController) DiscoverModels(ctx context.Context, accountID, clusterAddr string, req *proto.ModelDiscoveryRequest) (*proto.ModelDiscoveryResult, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "DiscoverModels", ctx, accountID, clusterAddr, req)
ret0, _ := ret[0].(*proto.ModelDiscoveryResult)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// DiscoverModels indicates an expected call of DiscoverModels.
func (mr *MockControllerMockRecorder) DiscoverModels(ctx, accountID, clusterAddr, req interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DiscoverModels", reflect.TypeOf((*MockController)(nil).DiscoverModels), ctx, accountID, clusterAddr, req)
}
// GetOIDCValidationConfig mocks base method.
func (m *MockController) GetOIDCValidationConfig() OIDCValidationConfig {
m.ctrl.T.Helper()

View File

@@ -24,13 +24,13 @@ import (
"github.com/netbirdio/netbird/encryption"
"github.com/netbirdio/netbird/formatter/hook"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs"
accesslogsmanager "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs/manager"
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
"github.com/netbirdio/netbird/management/server/activity"
activitystore "github.com/netbirdio/netbird/management/server/activity/store"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
nbcache "github.com/netbirdio/netbird/management/server/cache"
nbContext "github.com/netbirdio/netbird/management/server/context"
nbhttp "github.com/netbirdio/netbird/management/server/http"

View File

@@ -8,7 +8,6 @@ import (
"github.com/netbirdio/management-integrations/integrations"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
"github.com/netbirdio/netbird/management/internals/modules/peers"
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/domain/manager"
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
@@ -21,6 +20,7 @@ import (
recordsManager "github.com/netbirdio/netbird/management/internals/modules/zones/records/manager"
"github.com/netbirdio/netbird/management/server"
"github.com/netbirdio/netbird/management/server/account"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
"github.com/netbirdio/netbird/management/server/geolocation"
"github.com/netbirdio/netbird/management/server/groups"
"github.com/netbirdio/netbird/management/server/idp"

View File

@@ -0,0 +1,404 @@
package grpc
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
rpproxy "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
"github.com/netbirdio/netbird/shared/management/proto"
nbstatus "github.com/netbirdio/netbird/shared/management/status"
)
type modelDiscoveryCall struct {
result *proto.ModelDiscoveryResult
err error
}
func newModelDiscoveryTestServer(t *testing.T) (*ProxyServiceServer, *testProxyController) {
t.Helper()
controller := newTestProxyController()
server := &ProxyServiceServer{
proxyController: controller,
modelDiscoveryPending: make(map[string]*pendingModelDiscovery),
}
return server, controller
}
func addModelDiscoveryTestConnection(
t *testing.T,
server *ProxyServiceServer,
controller *testProxyController,
proxyID, sessionID, cluster string,
accountID *string,
supportsDiscovery bool,
) (*proxyConnection, context.CancelFunc) {
t.Helper()
ctx, cancel := context.WithCancel(context.Background())
ready := make(chan struct{})
close(ready)
conn := &proxyConnection{
proxyID: proxyID,
sessionID: sessionID,
address: cluster,
accountID: accountID,
capabilities: &proto.ProxyCapabilities{
SupportsModelDiscovery: &supportsDiscovery,
},
syncStream: &syncRecordingStream{},
modelDiscoveryChan: make(chan *proto.ModelDiscoveryRequest, modelDiscoveryQueueSize),
ready: ready,
ctx: ctx,
cancel: cancel,
}
server.connectedProxies.Store(proxyID, conn)
require.NoError(t, controller.RegisterProxyToCluster(context.Background(), cluster, proxyID))
t.Cleanup(func() {
cancel()
server.connectedProxies.Delete(proxyID)
_ = controller.UnregisterProxyFromCluster(context.Background(), cluster, proxyID)
})
return conn, cancel
}
func callDiscoverModels(
server *ProxyServiceServer,
ctx context.Context,
accountID, cluster string,
req *proto.ModelDiscoveryRequest,
) <-chan modelDiscoveryCall {
done := make(chan modelDiscoveryCall, 1)
go func() {
result, err := server.DiscoverModels(ctx, accountID, cluster, req)
done <- modelDiscoveryCall{result: result, err: err}
}()
return done
}
func TestDiscoverModels_CorrelatesResultAndCopiesSensitiveRequest(t *testing.T) {
server, controller := newModelDiscoveryTestServer(t)
conn, _ := addModelDiscoveryTestConnection(
t, server, controller, "proxy-a", "session-a", "cluster.example.com", nil, true,
)
callerReq := &proto.ModelDiscoveryRequest{
RequestId: "caller-controlled",
UpstreamUrl: "http://ollama.internal:11434",
AuthHeaderName: "Authorization",
AuthHeaderValue: "Bearer secret",
SkipTlsVerify: true,
OllamaFallback: true,
}
done := callDiscoverModels(server, context.Background(), "account-a", "cluster.example.com", callerReq)
wireReq := <-conn.modelDiscoveryChan
assert.NotEmpty(t, wireReq.GetRequestId())
assert.NotEqual(t, callerReq.GetRequestId(), wireReq.GetRequestId(), "management must own correlation IDs")
assert.Equal(t, callerReq.GetUpstreamUrl(), wireReq.GetUpstreamUrl())
assert.Equal(t, callerReq.GetAuthHeaderName(), wireReq.GetAuthHeaderName())
assert.Equal(t, callerReq.GetAuthHeaderValue(), wireReq.GetAuthHeaderValue())
assert.Equal(t, callerReq.GetSkipTlsVerify(), wireReq.GetSkipTlsVerify())
assert.Equal(t, callerReq.GetOllamaFallback(), wireReq.GetOllamaFallback())
assert.Equal(t, "caller-controlled", callerReq.GetRequestId(), "caller request must not be mutated")
want := &proto.ModelDiscoveryResult{
RequestId: wireReq.GetRequestId(),
Models: []*proto.ModelDiscoveryModel{
{Id: "llama3.2:latest", Label: "llama3.2:latest"},
},
Source: "openai_v1_models",
}
server.completeModelDiscovery(conn, want)
got := <-done
require.NoError(t, got.err)
assert.Equal(t, want, got.result)
}
func TestDiscoverModels_IgnoresMismatchedCorrelation(t *testing.T) {
server, controller := newModelDiscoveryTestServer(t)
conn, _ := addModelDiscoveryTestConnection(
t, server, controller, "proxy-a", "session-a", "cluster.example.com", nil, true,
)
done := callDiscoverModels(server, context.Background(), "account-a", "cluster.example.com", &proto.ModelDiscoveryRequest{
UpstreamUrl: "http://ollama.internal:11434",
})
wireReq := <-conn.modelDiscoveryChan
server.completeModelDiscovery(conn, &proto.ModelDiscoveryResult{RequestId: "wrong-request"})
select {
case got := <-done:
t.Fatalf("mismatched request completed discovery: %+v", got)
default:
}
wrongSession := *conn
wrongSession.sessionID = "session-b"
server.completeModelDiscovery(&wrongSession, &proto.ModelDiscoveryResult{RequestId: wireReq.GetRequestId()})
select {
case got := <-done:
t.Fatalf("mismatched proxy session completed discovery: %+v", got)
default:
}
server.completeModelDiscovery(conn, &proto.ModelDiscoveryResult{RequestId: wireReq.GetRequestId()})
got := <-done
require.NoError(t, got.err)
require.NotNil(t, got.result)
}
func TestDiscoverModels_ReturnsProxyReportedErrorForManagerContext(t *testing.T) {
server, controller := newModelDiscoveryTestServer(t)
conn, _ := addModelDiscoveryTestConnection(
t, server, controller, "proxy-a", "session-a", "cluster.example.com", nil, true,
)
done := callDiscoverModels(server, context.Background(), "account-a", "cluster.example.com", &proto.ModelDiscoveryRequest{
UpstreamUrl: "http://ollama.internal:11434",
})
wireReq := <-conn.modelDiscoveryChan
server.completeModelDiscovery(conn, &proto.ModelDiscoveryResult{
RequestId: wireReq.GetRequestId(),
Error: "upstream returned HTTP 401",
})
got := <-done
require.NoError(t, got.err)
require.NotNil(t, got.result)
assert.Equal(t, "upstream returned HTTP 401", got.result.GetError())
assert.Equal(t, wireReq.GetRequestId(), got.result.GetRequestId())
}
func TestDiscoverModels_FiltersCapabilityAndAccountScope(t *testing.T) {
tests := []struct {
name string
connectionAccount *string
requestAccount string
supported bool
}{
{
name: "capability absent",
requestAccount: "account-a",
supported: false,
},
{
name: "BYOP account mismatch",
connectionAccount: stringPointer("account-b"),
requestAccount: "account-a",
supported: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
server, controller := newModelDiscoveryTestServer(t)
addModelDiscoveryTestConnection(
t, server, controller, "proxy-a", "session-a", "cluster.example.com",
tt.connectionAccount, tt.supported,
)
result, err := server.DiscoverModels(
context.Background(),
tt.requestAccount,
"cluster.example.com",
&proto.ModelDiscoveryRequest{UpstreamUrl: "http://ollama.internal:11434"},
)
require.Nil(t, result)
requireStatusType(t, err, nbstatus.PreconditionFailed)
})
}
}
func TestDiscoverModels_RejectsStaleClusterMembershipAfterReconnect(t *testing.T) {
server, controller := newModelDiscoveryTestServer(t)
conn, _ := addModelDiscoveryTestConnection(
t, server, controller, "proxy-a", "new-session", "new-cluster.example.com", nil, true,
)
// Simulate the stale membership left behind when this proxy ID previously
// connected to another cluster and that superseded session skipped cleanup.
require.NoError(t, controller.RegisterProxyToCluster(
context.Background(),
"old-cluster.example.com",
"proxy-a",
))
result, err := server.DiscoverModels(
context.Background(),
"account-a",
"old-cluster.example.com",
&proto.ModelDiscoveryRequest{UpstreamUrl: "http://ollama.internal:11434"},
)
require.Nil(t, result)
requireStatusType(t, err, nbstatus.PreconditionFailed)
select {
case request := <-conn.modelDiscoveryChan:
t.Fatalf("sent credentials to connection in a different cluster: %+v", request)
default:
}
}
func TestDiscoverModels_CancelAndDisconnectCleanPendingRequest(t *testing.T) {
t.Run("caller cancellation", func(t *testing.T) {
server, controller := newModelDiscoveryTestServer(t)
conn, _ := addModelDiscoveryTestConnection(
t, server, controller, "proxy-a", "session-a", "cluster.example.com", nil, true,
)
ctx, cancel := context.WithCancel(context.Background())
done := callDiscoverModels(server, ctx, "account-a", "cluster.example.com", &proto.ModelDiscoveryRequest{
UpstreamUrl: "http://ollama.internal:11434",
})
<-conn.modelDiscoveryChan
cancel()
got := <-done
require.Nil(t, got.result)
requireStatusType(t, got.err, nbstatus.PreconditionFailed)
assert.Empty(t, server.modelDiscoveryPending)
})
t.Run("proxy disconnect", func(t *testing.T) {
server, controller := newModelDiscoveryTestServer(t)
conn, _ := addModelDiscoveryTestConnection(
t, server, controller, "proxy-a", "session-a", "cluster.example.com", nil, true,
)
done := callDiscoverModels(server, context.Background(), "account-a", "cluster.example.com", &proto.ModelDiscoveryRequest{
UpstreamUrl: "http://ollama.internal:11434",
})
<-conn.modelDiscoveryChan
server.failModelDiscoveries(conn, rpproxy.ErrModelDiscoveryUnavailable)
got := <-done
require.Nil(t, got.result)
requireStatusType(t, got.err, nbstatus.PreconditionFailed)
assert.Empty(t, server.modelDiscoveryPending)
})
}
func TestDiscoverModels_UsesNextCapableProxyWhenFirstQueueIsFull(t *testing.T) {
server, controller := newModelDiscoveryTestServer(t)
first, _ := addModelDiscoveryTestConnection(
t, server, controller, "proxy-a", "session-a", "cluster.example.com", nil, true,
)
second, _ := addModelDiscoveryTestConnection(
t, server, controller, "proxy-b", "session-b", "cluster.example.com", nil, true,
)
first.modelDiscoveryChan = make(chan *proto.ModelDiscoveryRequest, 1)
first.modelDiscoveryChan <- &proto.ModelDiscoveryRequest{RequestId: "occupied"}
done := callDiscoverModels(server, context.Background(), "account-a", "cluster.example.com", &proto.ModelDiscoveryRequest{
UpstreamUrl: "http://ollama.internal:11434",
})
wireReq := <-second.modelDiscoveryChan
server.completeModelDiscovery(second, &proto.ModelDiscoveryResult{RequestId: wireReq.GetRequestId()})
got := <-done
require.NoError(t, got.err)
require.NotNil(t, got.result)
}
func TestSenderSkipsCanceledQueuedModelDiscovery(t *testing.T) {
server, controller := newModelDiscoveryTestServer(t)
conn, _ := addModelDiscoveryTestConnection(
t, server, controller, "proxy-a", "session-a", "cluster.example.com", nil, true,
)
stream := conn.syncStream.(*syncRecordingStream)
ctx, cancel := context.WithCancel(context.Background())
done := callDiscoverModels(server, ctx, "account-a", "cluster.example.com", &proto.ModelDiscoveryRequest{
UpstreamUrl: "http://ollama.internal:11434",
})
require.Eventually(t, func() bool {
return len(conn.modelDiscoveryChan) == 1
}, time.Second, time.Millisecond)
cancel()
got := <-done
require.Nil(t, got.result)
requireStatusType(t, got.err, nbstatus.PreconditionFailed)
liveRequest := &proto.ModelDiscoveryRequest{RequestId: "live-request"}
livePending := &pendingModelDiscovery{
proxyID: conn.proxyID,
sessionID: conn.sessionID,
done: make(chan modelDiscoveryCompletion, 1),
}
server.registerModelDiscovery(liveRequest.GetRequestId(), livePending)
defer server.removeModelDiscovery(liveRequest.GetRequestId(), livePending)
conn.modelDiscoveryChan <- liveRequest
errCh := make(chan error, 1)
go server.sender(conn, errCh)
require.Eventually(t, func() bool {
stream.mu.Lock()
defer stream.mu.Unlock()
return len(stream.sent) == 1
}, time.Second, time.Millisecond)
stream.mu.Lock()
sent := append([]*proto.SyncMappingsResponse(nil), stream.sent...)
stream.mu.Unlock()
require.Len(t, sent, 1)
assert.Equal(t, "live-request", sent[0].GetModelDiscoveryRequest().GetRequestId())
}
func TestDrainRecv_DispatchesModelDiscoveryResult(t *testing.T) {
server, _ := newModelDiscoveryTestServer(t)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
conn := &proxyConnection{
proxyID: "proxy-a",
sessionID: "session-a",
ctx: ctx,
cancel: cancel,
}
pending := &pendingModelDiscovery{
proxyID: conn.proxyID,
sessionID: conn.sessionID,
done: make(chan modelDiscoveryCompletion, 1),
}
server.registerModelDiscovery("request-a", pending)
stream := &syncRecordingStream{
recvMsgs: []*proto.SyncMappingsRequest{
{
Msg: &proto.SyncMappingsRequest_ModelDiscoveryResult{
ModelDiscoveryResult: &proto.ModelDiscoveryResult{RequestId: "request-a"},
},
},
},
}
errCh := make(chan error, 1)
go server.drainRecv(conn, stream, errCh)
completion := <-pending.done
require.NoError(t, completion.err)
require.NotNil(t, completion.result)
assert.Equal(t, "request-a", completion.result.GetRequestId())
}
func TestSendModelDiscoveryRequest_UsesOutOfBandSyncField(t *testing.T) {
stream := &syncRecordingStream{}
conn := &proxyConnection{syncStream: stream}
req := &proto.ModelDiscoveryRequest{RequestId: "request-a"}
require.NoError(t, conn.sendModelDiscoveryRequest(req))
require.Len(t, stream.sent, 1)
assert.Empty(t, stream.sent[0].GetMapping())
assert.Equal(t, req, stream.sent[0].GetModelDiscoveryRequest())
}
func stringPointer(value string) *string {
return &value
}
func requireStatusType(t *testing.T, err error, want nbstatus.Type) {
t.Helper()
require.Error(t, err)
statusErr, ok := nbstatus.FromError(err)
require.True(t, ok)
assert.Equal(t, want, statusErr.Type())
}

View File

@@ -15,6 +15,7 @@ import (
"net/http"
"net/url"
"os"
"sort"
"strconv"
"strings"
"sync"
@@ -84,6 +85,10 @@ type ProxyServiceServer struct {
// Map of connected proxies: proxy_id -> proxy connection
connectedProxies sync.Map
// modelDiscoveryPending correlates out-of-band model-discovery results
// received on SyncMappings with the management request waiting for them.
modelDiscoveryMu sync.Mutex
modelDiscoveryPending map[string]*pendingModelDiscovery
// Manager for access logs
accessLogManager accesslogs.Manager
@@ -143,6 +148,8 @@ const defaultProxyTokenTTL = 5 * time.Minute
const defaultSnapshotBatchSize = 500
const modelDiscoveryQueueSize = 16
func snapshotBatchSizeFromEnv() int {
if v := os.Getenv("NB_PROXY_SNAPSHOT_BATCH_SIZE"); v != "" {
if n, err := strconv.Atoi(v); err == nil && n > 0 {
@@ -171,10 +178,26 @@ type proxyConnection struct {
stream proto.ProxyService_GetMappingUpdateServer
// syncStream is set when the proxy connected via SyncMappings.
// When non-nil, the sender goroutine uses this instead of stream.
syncStream proto.ProxyService_SyncMappingsServer
sendChan chan *proto.GetMappingUpdateResponse
ctx context.Context
cancel context.CancelFunc
syncStream proto.ProxyService_SyncMappingsServer
sendChan chan *proto.GetMappingUpdateResponse
modelDiscoveryChan chan *proto.ModelDiscoveryRequest
// ready closes after the initial snapshot completes and the sender
// goroutine is about to start. Discovery never competes with snapshot
// delivery on the stream.
ready chan struct{}
ctx context.Context
cancel context.CancelFunc
}
type modelDiscoveryCompletion struct {
result *proto.ModelDiscoveryResult
err error
}
type pendingModelDiscovery struct {
proxyID string
sessionID string
done chan modelDiscoveryCompletion
}
func enforceAccountScope(ctx context.Context, requestAccountID string) error {
@@ -192,17 +215,18 @@ func enforceAccountScope(ctx context.Context, requestAccountID string) error {
func NewProxyServiceServer(accessLogMgr accesslogs.Manager, tokenStore *OneTimeTokenStore, pkceStore *PKCEVerifierStore, oidcConfig ProxyOIDCConfig, peersManager peers.Manager, usersManager users.Manager, idpManager idp.Manager, proxyMgr proxy.Manager, tokenChecker ProxyTokenChecker) *ProxyServiceServer {
ctx, cancel := context.WithCancel(context.Background())
s := &ProxyServiceServer{
accessLogManager: accessLogMgr,
oidcConfig: oidcConfig,
tokenStore: tokenStore,
pkceVerifierStore: pkceStore,
peersManager: peersManager,
usersManager: usersManager,
idpManager: idpManager,
proxyManager: proxyMgr,
tokenChecker: tokenChecker,
snapshotBatchSize: snapshotBatchSizeFromEnv(),
cancel: cancel,
accessLogManager: accessLogMgr,
oidcConfig: oidcConfig,
tokenStore: tokenStore,
pkceVerifierStore: pkceStore,
peersManager: peersManager,
usersManager: usersManager,
idpManager: idpManager,
proxyManager: proxyMgr,
tokenChecker: tokenChecker,
snapshotBatchSize: snapshotBatchSizeFromEnv(),
modelDiscoveryPending: make(map[string]*pendingModelDiscovery),
cancel: cancel,
}
go s.cleanupStaleProxies(ctx)
return s
@@ -392,6 +416,7 @@ func (s *ProxyServiceServer) GetMappingUpdate(req *proto.GetMappingUpdateRequest
return fmt.Errorf("send snapshot to proxy %s: %w", params.proxyID, err)
}
close(conn.ready)
errChan := make(chan error, 2)
go s.sender(conn, errChan)
@@ -425,9 +450,10 @@ func (s *ProxyServiceServer) SyncMappings(stream proto.ProxyService_SyncMappings
return fmt.Errorf("send snapshot to proxy %s: %w", params.proxyID, err)
}
close(conn.ready)
errChan := make(chan error, 2)
go s.sender(conn, errChan)
go s.drainRecv(stream, errChan)
go s.drainRecv(conn, stream, errChan)
return s.serveProxyConnection(conn, proxyRecord, errChan, true)
}
@@ -496,6 +522,8 @@ func (s *ProxyServiceServer) registerProxyConnection(ctx context.Context, params
connSeed.tokenID = tokenID
connSeed.capabilities = params.capabilities
connSeed.sendChan = make(chan *proto.GetMappingUpdateResponse, 100)
connSeed.modelDiscoveryChan = make(chan *proto.ModelDiscoveryRequest, modelDiscoveryQueueSize)
connSeed.ready = make(chan struct{})
connSeed.ctx = connCtx
connSeed.cancel = cancel
@@ -543,10 +571,13 @@ func (s *ProxyServiceServer) supersedePriorConnection(proxyID, newSessionID stri
// cleanupFailedSnapshot removes the connection from the cluster and store
// after a snapshot send failure.
func (s *ProxyServiceServer) cleanupFailedSnapshot(ctx context.Context, conn *proxyConnection) {
s.failModelDiscoveries(conn, proxy.ErrModelDiscoveryUnavailable)
if s.connectedProxies.CompareAndDelete(conn.proxyID, conn) {
if err := s.proxyController.UnregisterProxyFromCluster(context.Background(), conn.address, conn.proxyID); err != nil {
log.WithContext(ctx).Debugf("cleanup after snapshot failure for proxy %s: %v", conn.proxyID, err)
}
} else {
s.unregisterSupersededProxyCluster(ctx, conn)
}
conn.cancel()
if err := s.proxyManager.Disconnect(context.Background(), conn.proxyID, conn.sessionID); err != nil {
@@ -554,15 +585,19 @@ func (s *ProxyServiceServer) cleanupFailedSnapshot(ctx context.Context, conn *pr
}
}
// drainRecv consumes and discards messages from a bidirectional stream.
// The proxy sends an ack for every incremental update; we don't need them
// after the snapshot phase. Recv errors are forwarded to errChan.
func (s *ProxyServiceServer) drainRecv(stream proto.ProxyService_SyncMappingsServer, errChan chan<- error) {
// drainRecv consumes post-snapshot messages from a bidirectional stream.
// Incremental mapping acks need no action; model-discovery results are
// dispatched to the correlated management caller.
func (s *ProxyServiceServer) drainRecv(conn *proxyConnection, stream proto.ProxyService_SyncMappingsServer, errChan chan<- error) {
for {
if _, err := stream.Recv(); err != nil {
msg, err := stream.Recv()
if err != nil {
errChan <- err
return
}
if result := msg.GetModelDiscoveryResult(); result != nil {
s.completeModelDiscovery(conn, result)
}
}
}
@@ -599,7 +634,9 @@ func (s *ProxyServiceServer) serveProxyConnection(conn *proxyConnection, proxyRe
// disconnectProxy removes the connection from cluster and store, unless it
// has already been superseded by a newer connection.
func (s *ProxyServiceServer) disconnectProxy(conn *proxyConnection) {
s.failModelDiscoveries(conn, proxy.ErrModelDiscoveryUnavailable)
if !s.connectedProxies.CompareAndDelete(conn.proxyID, conn) {
s.unregisterSupersededProxyCluster(context.Background(), conn)
log.Infof("Proxy %s session %s: skipping cleanup, superseded by new connection", conn.proxyID, conn.sessionID)
conn.cancel()
return
@@ -616,6 +653,26 @@ func (s *ProxyServiceServer) disconnectProxy(conn *proxyConnection) {
log.Infof("Proxy %s session %s disconnected", conn.proxyID, conn.sessionID)
}
// unregisterSupersededProxyCluster removes membership owned only by an old
// session after the same proxy ID reconnects at a different address. When the
// address is unchanged, the membership is shared with the replacement and
// must remain registered.
func (s *ProxyServiceServer) unregisterSupersededProxyCluster(ctx context.Context, conn *proxyConnection) {
if conn == nil || s.proxyController == nil {
return
}
currentValue, ok := s.connectedProxies.Load(conn.proxyID)
if ok {
current := currentValue.(*proxyConnection)
if current == conn || current.address == conn.address {
return
}
}
if err := s.proxyController.UnregisterProxyFromCluster(ctx, conn.address, conn.proxyID); err != nil {
log.WithContext(ctx).Debugf("cleanup superseded cluster membership for proxy %s: %v", conn.proxyID, err)
}
}
// sendSnapshotSync sends the initial snapshot with back-pressure: it sends
// one batch, then waits for the proxy to ack before sending the next.
func (s *ProxyServiceServer) sendSnapshotSync(ctx context.Context, conn *proxyConnection, stream proto.ProxyService_SyncMappingsServer) error {
@@ -852,6 +909,20 @@ func (s *ProxyServiceServer) sender(conn *proxyConnection, errChan chan<- error)
return
}
log.WithContext(conn.ctx).Tracef("Send response to proxy %s", conn.proxyID)
case req := <-conn.modelDiscoveryChan:
// The API caller may have timed out while this request waited
// behind other stream work. Pending correlation is removed on
// cancellation, so skip stale probes before they reach the
// proxy and expose a stored credential unnecessarily.
if !s.isModelDiscoveryPendingFor(conn, req.GetRequestId()) {
continue
}
if err := conn.sendModelDiscoveryRequest(req); err != nil {
errChan <- err
log.WithContext(conn.ctx).Tracef("Failed to send model discovery request to proxy %s: %v", conn.proxyID, err)
return
}
log.WithContext(conn.ctx).Tracef("Sent model discovery request to proxy %s", conn.proxyID)
case <-conn.ctx.Done():
return
}
@@ -869,6 +940,15 @@ func (conn *proxyConnection) sendResponse(resp *proto.GetMappingUpdateResponse)
return conn.stream.Send(resp)
}
func (conn *proxyConnection) sendModelDiscoveryRequest(req *proto.ModelDiscoveryRequest) error {
if conn.syncStream == nil {
return proxy.ErrModelDiscoveryUnavailable
}
return conn.syncStream.Send(&proto.SyncMappingsResponse{
ModelDiscoveryRequest: req,
})
}
// SendAccessLog processes access log from proxy
func (s *ProxyServiceServer) SendAccessLog(ctx context.Context, req *proto.SendAccessLogRequest) (*proto.SendAccessLogResponse, error) {
accessLog := req.GetLog()
@@ -1017,6 +1097,181 @@ func (s *ProxyServiceServer) GetConnectedProxyURLs() []string {
return urls
}
// DiscoverModels sends one correlated model-discovery request to a capable
// proxy in clusterAddr and waits for its result. Only the stored connection's
// cluster/account scope is used for routing; no routing identity is carried in
// the request delivered to the proxy.
func (s *ProxyServiceServer) DiscoverModels(ctx context.Context, accountID, clusterAddr string, req *proto.ModelDiscoveryRequest) (*proto.ModelDiscoveryResult, error) {
if req == nil {
return nil, nbstatus.Errorf(nbstatus.InvalidArgument, "model discovery request is required")
}
if strings.TrimSpace(accountID) == "" {
return nil, nbstatus.Errorf(nbstatus.InvalidArgument, "account ID is required for model discovery")
}
if strings.TrimSpace(clusterAddr) == "" {
return nil, nbstatus.Errorf(nbstatus.PreconditionFailed, "agent network proxy cluster is not configured")
}
if s.proxyController == nil {
return nil, nbstatus.Errorf(nbstatus.PreconditionFailed, "model discovery is unavailable for the configured proxy cluster")
}
proxyIDs := s.proxyController.GetProxiesForCluster(clusterAddr)
sort.Strings(proxyIDs)
request := &proto.ModelDiscoveryRequest{
RequestId: uuid.NewString(),
UpstreamUrl: req.GetUpstreamUrl(),
AuthHeaderName: req.GetAuthHeaderName(),
AuthHeaderValue: req.GetAuthHeaderValue(),
SkipTlsVerify: req.GetSkipTlsVerify(),
OllamaFallback: req.GetOllamaFallback(),
}
for _, proxyID := range proxyIDs {
connVal, ok := s.connectedProxies.Load(proxyID)
if !ok {
continue
}
conn := connVal.(*proxyConnection)
// Cluster membership can briefly retain a proxy ID from a superseded
// session. The live connection is authoritative: never send provider
// credentials unless it is connected to the exact requested cluster.
if conn.address != clusterAddr {
continue
}
if !modelDiscoveryCapable(conn, accountID) {
continue
}
pending := &pendingModelDiscovery{
proxyID: conn.proxyID,
sessionID: conn.sessionID,
done: make(chan modelDiscoveryCompletion, 1),
}
s.registerModelDiscovery(request.GetRequestId(), pending)
queued := false
select {
case conn.modelDiscoveryChan <- request:
queued = true
case <-ctx.Done():
s.removeModelDiscovery(request.GetRequestId(), pending)
return nil, modelDiscoveryContextError(ctx.Err())
default:
s.removeModelDiscovery(request.GetRequestId(), pending)
}
if !queued {
continue
}
defer s.removeModelDiscovery(request.GetRequestId(), pending)
select {
case completion := <-pending.done:
if completion.err != nil {
return nil, nbstatus.Errorf(nbstatus.PreconditionFailed, "model discovery proxy disconnected")
}
return completion.result, nil
case <-ctx.Done():
return nil, modelDiscoveryContextError(ctx.Err())
case <-conn.ctx.Done():
return nil, nbstatus.Errorf(nbstatus.PreconditionFailed, "model discovery proxy disconnected")
}
}
return nil, nbstatus.Errorf(nbstatus.PreconditionFailed, "no connected proxy in the configured cluster supports model discovery")
}
func modelDiscoveryCapable(conn *proxyConnection, accountID string) bool {
if conn == nil || conn.syncStream == nil || conn.modelDiscoveryChan == nil || conn.ready == nil {
return false
}
if conn.ctx == nil || conn.ctx.Err() != nil {
return false
}
if conn.accountID != nil && *conn.accountID != accountID {
return false
}
if conn.capabilities == nil || !conn.capabilities.GetSupportsModelDiscovery() {
return false
}
select {
case <-conn.ready:
return true
default:
return false
}
}
func modelDiscoveryContextError(err error) error {
if errors.Is(err, context.DeadlineExceeded) {
return nbstatus.Errorf(nbstatus.PreconditionFailed, "model discovery timed out")
}
return nbstatus.Errorf(nbstatus.PreconditionFailed, "model discovery was canceled")
}
func (s *ProxyServiceServer) registerModelDiscovery(requestID string, pending *pendingModelDiscovery) {
s.modelDiscoveryMu.Lock()
defer s.modelDiscoveryMu.Unlock()
if s.modelDiscoveryPending == nil {
s.modelDiscoveryPending = make(map[string]*pendingModelDiscovery)
}
s.modelDiscoveryPending[requestID] = pending
}
func (s *ProxyServiceServer) removeModelDiscovery(requestID string, pending *pendingModelDiscovery) {
s.modelDiscoveryMu.Lock()
defer s.modelDiscoveryMu.Unlock()
if s.modelDiscoveryPending[requestID] == pending {
delete(s.modelDiscoveryPending, requestID)
}
}
func (s *ProxyServiceServer) isModelDiscoveryPendingFor(conn *proxyConnection, requestID string) bool {
if conn == nil || requestID == "" {
return false
}
s.modelDiscoveryMu.Lock()
defer s.modelDiscoveryMu.Unlock()
pending := s.modelDiscoveryPending[requestID]
return pending != nil && pending.proxyID == conn.proxyID && pending.sessionID == conn.sessionID
}
func (s *ProxyServiceServer) completeModelDiscovery(conn *proxyConnection, result *proto.ModelDiscoveryResult) {
if conn == nil || result == nil || result.GetRequestId() == "" {
return
}
s.modelDiscoveryMu.Lock()
pending := s.modelDiscoveryPending[result.GetRequestId()]
s.modelDiscoveryMu.Unlock()
if pending == nil || pending.proxyID != conn.proxyID || pending.sessionID != conn.sessionID {
return
}
select {
case pending.done <- modelDiscoveryCompletion{result: result}:
default:
}
}
func (s *ProxyServiceServer) failModelDiscoveries(conn *proxyConnection, err error) {
if conn == nil {
return
}
s.modelDiscoveryMu.Lock()
pending := make([]*pendingModelDiscovery, 0)
for _, item := range s.modelDiscoveryPending {
if item.proxyID == conn.proxyID && item.sessionID == conn.sessionID {
pending = append(pending, item)
}
}
s.modelDiscoveryMu.Unlock()
for _, item := range pending {
select {
case item.done <- modelDiscoveryCompletion{err: err}:
default:
}
}
}
// SendServiceUpdateToCluster sends a service update to all proxy servers in a specific cluster.
// If clusterAddr is empty, broadcasts to all connected proxy servers (backward compatibility).
// For create/update operations a unique one-time auth token is generated per
@@ -1049,6 +1304,12 @@ func (s *ProxyServiceServer) SendServiceUpdateToCluster(ctx context.Context, upd
continue
}
conn := connVal.(*proxyConnection)
// Membership can retain this proxy ID from a superseded session in a
// different cluster. The live connection address is authoritative,
// especially because Agent Network mappings can contain credentials.
if conn.address != clusterAddr {
continue
}
if conn.accountID != nil && update.AccountId != "" && *conn.accountID != update.AccountId {
continue
}

View File

@@ -42,6 +42,10 @@ func newTestProxyController() *testProxyController {
func (c *testProxyController) SendServiceUpdateToCluster(_ context.Context, _ string, _ *proto.ProxyMapping, _ string) {
}
func (c *testProxyController) DiscoverModels(_ context.Context, _, _ string, _ *proto.ModelDiscoveryRequest) (*proto.ModelDiscoveryResult, error) {
return nil, proxy.ErrModelDiscoveryUnavailable
}
func (c *testProxyController) GetOIDCValidationConfig() proxy.OIDCValidationConfig {
return proxy.OIDCValidationConfig{}
}
@@ -218,6 +222,65 @@ func TestSendServiceUpdateToCluster_DeleteNoToken(t *testing.T) {
assert.Empty(t, msg2.AuthToken)
}
func TestSendServiceUpdateToCluster_RejectsStaleClusterMembership(t *testing.T) {
ctx := context.Background()
s := &ProxyServiceServer{
tokenStore: NewOneTimeTokenStore(ctx, testCacheStore(t)),
}
controller := newTestProxyController()
s.SetProxyController(controller)
ch := registerFakeProxy(s, "proxy-a", "new-cluster.example.com")
require.NoError(t, controller.RegisterProxyToCluster(ctx, "old-cluster.example.com", "proxy-a"))
s.SendServiceUpdateToCluster(ctx, &proto.ProxyMapping{
Type: proto.ProxyMappingUpdateType_UPDATE_TYPE_CREATED,
Id: "agent-network-provider-1",
AccountId: "account-1",
Domain: "agent.example.com",
Path: []*proto.PathMapping{
{Path: "/", Target: "http://ollama.internal:11434/"},
},
}, "old-cluster.example.com")
assert.True(t, drainEmpty(ch), "stale membership must not route a mapping to a different live cluster")
}
func TestDisconnectProxy_RemovesSupersededMembershipFromOldCluster(t *testing.T) {
controller := newTestProxyController()
server := &ProxyServiceServer{proxyController: controller}
oldCtx, cancelOld := context.WithCancel(context.Background())
newCtx, cancelNew := context.WithCancel(context.Background())
t.Cleanup(cancelOld)
t.Cleanup(cancelNew)
oldConnection := &proxyConnection{
proxyID: "proxy-a",
sessionID: "old-session",
address: "old-cluster.example.com",
ctx: oldCtx,
cancel: cancelOld,
}
newConnection := &proxyConnection{
proxyID: "proxy-a",
sessionID: "new-session",
address: "new-cluster.example.com",
ctx: newCtx,
cancel: cancelNew,
}
require.NoError(t, controller.RegisterProxyToCluster(context.Background(), oldConnection.address, oldConnection.proxyID))
require.NoError(t, controller.RegisterProxyToCluster(context.Background(), newConnection.address, newConnection.proxyID))
server.connectedProxies.Store(newConnection.proxyID, newConnection)
server.disconnectProxy(oldConnection)
assert.Empty(t, controller.GetProxiesForCluster(oldConnection.address))
assert.Equal(t, []string{"proxy-a"}, controller.GetProxiesForCluster(newConnection.address))
current, ok := server.connectedProxies.Load(newConnection.proxyID)
require.True(t, ok)
assert.Same(t, newConnection, current)
}
func TestSendServiceUpdate_UniqueTokensPerProxy(t *testing.T) {
ctx := context.Background()
tokenStore := NewOneTimeTokenStore(ctx, testCacheStore(t))

View File

@@ -8,8 +8,8 @@ import (
"fmt"
"io"
"net/http"
"net/http/httptest"
"net/netip"
"net/http/httptest"
"strings"
"testing"

View File

@@ -191,3 +191,4 @@ func TestPolicyRuleEqual_EmptySlices(t *testing.T) {
}
assert.True(t, a.Equal(b))
}

View File

@@ -27,12 +27,12 @@ func TestOpenAIDetectFromURL(t *testing.T) {
// catches Cloudflare's OpenAI direct path
// (/v1/{account}/{gateway}/openai/chat/completions) and
// compat path (/v1/{account}/{gateway}/compat/chat/completions).
"/v1/{account}/{gateway}/openai/chat/completions": true,
"/v1/{account}/{gateway}/compat/chat/completions": true,
"/chat/completions": true,
"/v1/messages": false,
"/healthz": false,
"": false,
"/v1/{account}/{gateway}/openai/chat/completions": true,
"/v1/{account}/{gateway}/compat/chat/completions": true,
"/chat/completions": true,
"/v1/messages": false,
"/healthz": false,
"": false,
}
for path, want := range cases {
assert.Equal(t, want, p.DetectFromURL(path), "DetectFromURL(%q)", path)

View File

@@ -28,12 +28,12 @@ func TestDefaultTable_FirstPartyModelCoverage(t *testing.T) {
"ministral-8b-latest", "mistral-embed",
},
"anthropic": {
"claude-fable-5", "claude-opus-5", "claude-opus-4-8", "claude-opus-4-7", "claude-opus-4-6",
"claude-fable-5", "claude-opus-4-8", "claude-opus-4-7", "claude-opus-4-6",
"claude-opus-4-1", "claude-sonnet-4-6", "claude-sonnet-4-5", "claude-haiku-4-5",
},
// bedrock keys are the normalized ids the request parser emits.
"bedrock": {
"anthropic.claude-opus-5", "anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6",
"anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6",
"anthropic.claude-opus-4-1", "anthropic.claude-sonnet-4-6", "anthropic.claude-sonnet-4-5",
"anthropic.claude-haiku-4-5", "meta.llama3-3-70b-instruct",
"amazon.nova-pro", "amazon.nova-lite", "amazon.nova-micro", "amazon.nova-2-lite",

View File

@@ -180,11 +180,6 @@ anthropic:
output_per_1k: 0.050
cache_read_per_1k: 0.001
cache_creation_per_1k: 0.0125
claude-opus-5:
input_per_1k: 0.005
output_per_1k: 0.025
cache_read_per_1k: 0.0005
cache_creation_per_1k: 0.00625
claude-opus-4-8:
input_per_1k: 0.005
output_per_1k: 0.025
@@ -241,11 +236,6 @@ bedrock:
# eu.anthropic.claude-sonnet-4-5-20250929-v1:0 -> anthropic.claude-sonnet-4-5.
# Anthropic-on-Bedrock keeps the additive cache buckets (read ≈0.1x input,
# write ≈1.25x input); Nova / Llama report no cache, so cost is input+output.
anthropic.claude-opus-5:
input_per_1k: 0.005
output_per_1k: 0.025
cache_read_per_1k: 0.0005
cache_creation_per_1k: 0.00625
anthropic.claude-opus-4-8:
input_per_1k: 0.005
output_per_1k: 0.025

View File

@@ -14,9 +14,9 @@ import (
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/test/bufconn"
mgmtgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
mgmtgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
"github.com/netbirdio/netbird/management/server/store"
nbtypes "github.com/netbirdio/netbird/management/server/types"
"github.com/netbirdio/netbird/proxy/internal/middleware"

View File

@@ -0,0 +1,337 @@
// Package modeldiscovery fetches and normalizes model catalogs from
// Agent Network provider endpoints. Discovery always runs on the proxy so
// it observes the same network path as inference traffic.
package modeldiscovery
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"path"
"sort"
"strings"
"time"
"unicode"
"unicode/utf8"
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/proxy/internal/roundtrip"
)
const (
SourceOpenAIV1Models = "openai_v1_models"
SourceOllamaAPITags = "ollama_api_tags"
defaultTimeout = 5 * time.Second
maxResponseBytes = 1 << 20 // 1 MiB, after HTTP decompression.
maxUpstreamURLBytes = 4096
maxHeaderNameBytes = 256
maxHeaderValueBytes = 64 << 10
maxModels = 500
maxModelIDBytes = 512
)
// Request contains the provider-owned values management resolved from the
// persisted provider record. Callers must not populate these fields from a
// dashboard-supplied URL or credential.
type Request struct {
UpstreamURL string
AuthHeaderName string
AuthHeaderValue string
SkipTLSVerify bool
AllowOllamaFallback bool
}
// Model is the deliberately small response surface returned to management.
// Arbitrary fields supplied by an upstream never cross the control channel.
type Model struct {
ID string
Label string
}
// Result is a normalized model catalog and the endpoint shape that supplied
// it.
type Result struct {
Models []Model
Source string
}
// Discoverer owns the HTTP client used for provider probes.
type Discoverer struct {
client *http.Client
timeout time.Duration
bodyLimit int64
}
// New constructs a direct-upstream discoverer. The transport is the same
// host-network transport family used by Agent Network inference routes.
func New(logger *log.Logger) *Discoverer {
return newWithTransport(roundtrip.NewDirectOnly(logger))
}
func newWithTransport(transport http.RoundTripper) *Discoverer {
return &Discoverer{
client: &http.Client{
Transport: transport,
// Redirects could move a credentialed request away from the
// persisted provider origin. Discovery never follows them.
CheckRedirect: func(_ *http.Request, _ []*http.Request) error {
return http.ErrUseLastResponse
},
},
timeout: defaultTimeout,
bodyLimit: maxResponseBytes,
}
}
// Discover queries the OpenAI-compatible model-list endpoint. Ollama's native
// tags endpoint is attempted only when explicitly enabled and the primary
// endpoint reports that the route does not exist.
func (d *Discoverer) Discover(ctx context.Context, in Request) (Result, error) {
if d == nil || d.client == nil {
return Result{}, errors.New("model discovery client is unavailable")
}
if err := validateRequest(in); err != nil {
return Result{}, err
}
timeout := d.timeout
if timeout <= 0 || timeout > defaultTimeout {
timeout = defaultTimeout
}
probeCtx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
models, err := d.fetchOpenAIModels(probeCtx, in)
if err == nil {
return Result{Models: models, Source: SourceOpenAIV1Models}, nil
}
if !in.AllowOllamaFallback || !isMissingEndpoint(err) {
return Result{}, err
}
models, err = d.fetchOllamaTags(probeCtx, in)
if err != nil {
return Result{}, err
}
return Result{Models: models, Source: SourceOllamaAPITags}, nil
}
func validateRequest(in Request) error {
rawURL := strings.TrimSpace(in.UpstreamURL)
if rawURL == "" {
return errors.New("model discovery upstream URL is required")
}
if len(rawURL) > maxUpstreamURLBytes {
return errors.New("model discovery upstream URL is too long")
}
if len(in.AuthHeaderName) > maxHeaderNameBytes || len(in.AuthHeaderValue) > maxHeaderValueBytes {
return errors.New("model discovery authentication header is too large")
}
if (in.AuthHeaderName == "") != (in.AuthHeaderValue == "") {
return errors.New("model discovery authentication header is incomplete")
}
if strings.ContainsAny(in.AuthHeaderName, "\r\n") || strings.ContainsAny(in.AuthHeaderValue, "\r\n") {
return errors.New("model discovery authentication header is invalid")
}
if in.AuthHeaderName != "" && !strings.EqualFold(in.AuthHeaderName, "Authorization") {
return errors.New("model discovery authentication header is unsupported")
}
return nil
}
func (d *Discoverer) fetchOpenAIModels(ctx context.Context, in Request) ([]Model, error) {
body, err := d.fetch(ctx, in, "v1/models")
if err != nil {
return nil, err
}
var payload struct {
Data *[]struct {
ID string `json:"id"`
} `json:"data"`
}
if err := json.Unmarshal(body, &payload); err != nil || payload.Data == nil {
return nil, errors.New("upstream returned invalid OpenAI model-list JSON")
}
ids := make([]string, 0, len(*payload.Data))
for _, model := range *payload.Data {
ids = append(ids, model.ID)
}
return normalize(ids)
}
func (d *Discoverer) fetchOllamaTags(ctx context.Context, in Request) ([]Model, error) {
body, err := d.fetch(ctx, in, "api/tags")
if err != nil {
return nil, err
}
var payload struct {
Models *[]struct {
Name string `json:"name"`
Model string `json:"model"`
} `json:"models"`
}
if err := json.Unmarshal(body, &payload); err != nil || payload.Models == nil {
return nil, errors.New("upstream returned invalid Ollama tags JSON")
}
ids := make([]string, 0, len(*payload.Models))
for _, model := range *payload.Models {
id := model.Model
if strings.TrimSpace(id) == "" {
id = model.Name
}
ids = append(ids, id)
}
return normalize(ids)
}
func (d *Discoverer) fetch(ctx context.Context, in Request, endpointPath string) ([]byte, error) {
endpoint, err := buildEndpointURL(in.UpstreamURL, endpointPath)
if err != nil {
return nil, err
}
reqCtx := roundtrip.WithDirectUpstream(ctx)
if in.SkipTLSVerify {
reqCtx = roundtrip.WithSkipTLSVerify(reqCtx)
}
req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, endpoint.String(), nil)
if err != nil {
return nil, errors.New("create model discovery request")
}
req.Header.Set("Accept", "application/json")
if in.AuthHeaderName != "" {
// Discovery is currently enabled only for Ollama-compatible providers.
// Canonicalizing the sole catalog-owned credential header keeps the
// control message from becoming a generic arbitrary-header primitive.
req.Header.Set("Authorization", in.AuthHeaderValue)
}
resp, err := d.client.Do(req)
if err != nil {
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
return nil, errors.New("model discovery timed out")
}
if errors.Is(ctx.Err(), context.Canceled) {
return nil, context.Canceled
}
// Deliberately omit the underlying error: net/http errors include the
// internal URL, which should not be reflected through the public API.
return nil, errors.New("model discovery request failed")
}
defer func() {
_ = resp.Body.Close()
}()
if resp.StatusCode != http.StatusOK {
return nil, &upstreamStatusError{statusCode: resp.StatusCode}
}
limit := d.bodyLimit
if limit <= 0 || limit > maxResponseBytes {
limit = maxResponseBytes
}
if resp.ContentLength > limit {
return nil, errors.New("model discovery response is too large")
}
body, err := io.ReadAll(io.LimitReader(resp.Body, limit+1))
if err != nil {
return nil, errors.New("read model discovery response")
}
if int64(len(body)) > limit {
return nil, errors.New("model discovery response is too large")
}
return body, nil
}
func buildEndpointURL(rawURL, endpointPath string) (*url.URL, error) {
parsed, err := url.Parse(strings.TrimSpace(rawURL))
if err != nil || parsed.Host == "" || parsed.Hostname() == "" || parsed.Opaque != "" {
return nil, errors.New("model discovery upstream URL is invalid")
}
switch strings.ToLower(parsed.Scheme) {
case "http":
parsed.Scheme = "http"
case "https":
parsed.Scheme = "https"
default:
return nil, errors.New("model discovery upstream URL must use http or https")
}
if parsed.User != nil {
return nil, errors.New("model discovery upstream URL must not contain credentials")
}
// Match Agent Network routing semantics: the static discovery path is
// appended to any persisted base path. Queries and fragments on a provider
// URL are not forwarded to inference and are likewise excluded here.
parsed.Path = "/" + strings.TrimPrefix(path.Join(parsed.Path, endpointPath), "/")
parsed.RawPath = ""
parsed.RawQuery = ""
parsed.ForceQuery = false
parsed.Fragment = ""
return parsed, nil
}
func normalize(ids []string) ([]Model, error) {
if len(ids) > maxModels {
return nil, errors.New("upstream returned too many models")
}
seen := make(map[string]struct{}, len(ids))
normalized := make([]string, 0, len(ids))
for _, raw := range ids {
id := strings.TrimSpace(raw)
if !validModelID(id) {
continue
}
if _, ok := seen[id]; ok {
continue
}
seen[id] = struct{}{}
normalized = append(normalized, id)
}
sort.Strings(normalized)
models := make([]Model, 0, len(normalized))
for _, id := range normalized {
models = append(models, Model{ID: id, Label: id})
}
return models, nil
}
func validModelID(id string) bool {
if id == "" || len(id) > maxModelIDBytes || !utf8.ValidString(id) {
return false
}
for _, r := range id {
if unicode.IsControl(r) {
return false
}
}
return true
}
type upstreamStatusError struct {
statusCode int
}
func (e *upstreamStatusError) Error() string {
return fmt.Sprintf("upstream returned HTTP %d", e.statusCode)
}
func isMissingEndpoint(err error) bool {
var statusErr *upstreamStatusError
if !errors.As(err, &statusErr) {
return false
}
return statusErr.statusCode == http.StatusNotFound || statusErr.statusCode == http.StatusMethodNotAllowed
}

View File

@@ -0,0 +1,315 @@
package modeldiscovery
import (
"context"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestDiscoverOpenAIModels(t *testing.T) {
t.Parallel()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, http.MethodGet, r.Method)
assert.Equal(t, "/gateway/v1/models", r.URL.Path)
assert.Empty(t, r.URL.RawQuery)
assert.Equal(t, "application/json", r.Header.Get("Accept"))
assert.Equal(t, "Bearer secret", r.Header.Get("Authorization"))
_, _ = io.WriteString(w, `{
"object": "list",
"data": [
{"id": " qwen2.5:latest ", "owned_by": "ignored"},
{"id": "llama3.2:latest"},
{"id": "llama3.2:latest"},
{"id": ""},
{"id": "bad\u0000id"}
]
}`)
}))
defer server.Close()
discoverer := newWithTransport(http.DefaultTransport)
result, err := discoverer.Discover(context.Background(), Request{
UpstreamURL: server.URL + "/gateway/?ignored=true#fragment",
AuthHeaderName: "authorization",
AuthHeaderValue: "Bearer secret",
})
require.NoError(t, err)
assert.Equal(t, SourceOpenAIV1Models, result.Source)
assert.Equal(t, []Model{
{ID: "llama3.2:latest", Label: "llama3.2:latest"},
{ID: "qwen2.5:latest", Label: "qwen2.5:latest"},
}, result.Models)
}
func TestDiscoverOllamaFallback(t *testing.T) {
t.Parallel()
var primaryCalls atomic.Int32
var fallbackCalls atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/v1/models":
primaryCalls.Add(1)
http.NotFound(w, r)
case "/api/tags":
fallbackCalls.Add(1)
_, _ = io.WriteString(w, `{
"models": [
{"name": "ignored-name", "model": "gemma3:latest"},
{"name": "llama3.2:latest", "model": ""}
]
}`)
default:
http.Error(w, "unexpected path", http.StatusInternalServerError)
}
}))
defer server.Close()
result, err := newWithTransport(http.DefaultTransport).Discover(context.Background(), Request{
UpstreamURL: server.URL,
AllowOllamaFallback: true,
})
require.NoError(t, err)
assert.Equal(t, int32(1), primaryCalls.Load())
assert.Equal(t, int32(1), fallbackCalls.Load())
assert.Equal(t, SourceOllamaAPITags, result.Source)
assert.Equal(t, []Model{
{ID: "gemma3:latest", Label: "gemma3:latest"},
{ID: "llama3.2:latest", Label: "llama3.2:latest"},
}, result.Models)
}
func TestDiscoverDoesNotFallbackOnAuthenticationFailure(t *testing.T) {
t.Parallel()
var fallbackCalls atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/api/tags" {
fallbackCalls.Add(1)
}
http.Error(w, "secret upstream body", http.StatusUnauthorized)
}))
defer server.Close()
_, err := newWithTransport(http.DefaultTransport).Discover(context.Background(), Request{
UpstreamURL: server.URL,
AllowOllamaFallback: true,
})
require.EqualError(t, err, "upstream returned HTTP 401")
assert.Zero(t, fallbackCalls.Load())
assert.NotContains(t, err.Error(), "secret upstream body")
assert.NotContains(t, err.Error(), server.URL)
}
func TestDiscoverDoesNotFollowRedirectsOrLeakAuth(t *testing.T) {
t.Parallel()
var redirectedCalls atomic.Int32
redirected := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
redirectedCalls.Add(1)
assert.Empty(t, r.Header.Get("Authorization"))
_, _ = io.WriteString(w, `{"data":[]}`)
}))
defer redirected.Close()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
http.Redirect(w, r, redirected.URL+"/captured", http.StatusFound)
}))
defer server.Close()
_, err := newWithTransport(http.DefaultTransport).Discover(context.Background(), Request{
UpstreamURL: server.URL,
AuthHeaderName: "Authorization",
AuthHeaderValue: "Bearer must-not-leak",
})
require.EqualError(t, err, "upstream returned HTTP 302")
assert.Zero(t, redirectedCalls.Load())
}
func TestDiscoverEnforcesResponseLimit(t *testing.T) {
t.Parallel()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
_, _ = io.WriteString(w, `{"data":[{"id":"llama3.2:latest"}]}`)
}))
defer server.Close()
discoverer := newWithTransport(http.DefaultTransport)
discoverer.bodyLimit = 16
_, err := discoverer.Discover(context.Background(), Request{UpstreamURL: server.URL})
require.EqualError(t, err, "model discovery response is too large")
}
func TestDiscoverEnforcesTimeout(t *testing.T) {
t.Parallel()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
select {
case <-r.Context().Done():
case <-time.After(time.Second):
_, _ = io.WriteString(w, `{"data":[]}`)
}
}))
defer server.Close()
discoverer := newWithTransport(http.DefaultTransport)
discoverer.timeout = 30 * time.Millisecond
_, err := discoverer.Discover(context.Background(), Request{UpstreamURL: server.URL})
require.EqualError(t, err, "model discovery timed out")
}
func TestDiscoverHonorsSkipTLSVerification(t *testing.T) {
t.Parallel()
server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
_, _ = io.WriteString(w, `{"data":[]}`)
}))
defer server.Close()
discoverer := New(nil)
_, err := discoverer.Discover(context.Background(), Request{UpstreamURL: server.URL})
require.EqualError(t, err, "model discovery request failed")
result, err := discoverer.Discover(context.Background(), Request{
UpstreamURL: server.URL,
SkipTLSVerify: true,
})
require.NoError(t, err)
assert.Empty(t, result.Models)
}
func TestDiscoverRejectsUnsafeRequestValues(t *testing.T) {
t.Parallel()
tests := []struct {
name string
request Request
wantErr string
}{
{
name: "unsupported scheme",
request: Request{UpstreamURL: "file:///etc/passwd"},
wantErr: "model discovery upstream URL is invalid",
},
{
name: "URL credentials",
request: Request{UpstreamURL: "http://user:pass@example.com"},
wantErr: "model discovery upstream URL must not contain credentials",
},
{
name: "missing hostname",
request: Request{UpstreamURL: "http://:11434"},
wantErr: "model discovery upstream URL is invalid",
},
{
name: "incomplete auth header",
request: Request{
UpstreamURL: "http://example.com",
AuthHeaderName: "Authorization",
},
wantErr: "model discovery authentication header is incomplete",
},
{
name: "header injection",
request: Request{
UpstreamURL: "http://example.com",
AuthHeaderName: "Authorization",
AuthHeaderValue: "Bearer safe\r\nX-Evil: yes",
},
wantErr: "model discovery authentication header is invalid",
},
{
name: "unsupported auth header",
request: Request{
UpstreamURL: "http://example.com",
AuthHeaderName: "X-API-Key",
AuthHeaderValue: "secret",
},
wantErr: "model discovery authentication header is unsupported",
},
{
name: "oversized URL",
request: Request{
UpstreamURL: "http://example.com/" + strings.Repeat("a", maxUpstreamURLBytes),
},
wantErr: "model discovery upstream URL is too long",
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
t.Parallel()
_, err := newWithTransport(http.DefaultTransport).Discover(context.Background(), test.request)
require.EqualError(t, err, test.wantErr)
})
}
}
func TestDiscoverRejectsInvalidPayloads(t *testing.T) {
t.Parallel()
tests := []struct {
name string
status int
body string
wantErr string
}{
{
name: "malformed JSON",
status: http.StatusOK,
body: `{"data":`,
wantErr: "upstream returned invalid OpenAI model-list JSON",
},
{
name: "missing data",
status: http.StatusOK,
body: `{}`,
wantErr: "upstream returned invalid OpenAI model-list JSON",
},
{
name: "too many models",
status: http.StatusOK,
body: modelListJSON(maxModels + 1),
wantErr: "upstream returned too many models",
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
t.Parallel()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(test.status)
_, _ = io.WriteString(w, test.body)
}))
defer server.Close()
_, err := newWithTransport(http.DefaultTransport).Discover(context.Background(), Request{
UpstreamURL: server.URL,
})
require.EqualError(t, err, test.wantErr)
})
}
}
func modelListJSON(count int) string {
var body strings.Builder
body.WriteString(`{"data":[`)
for i := range count {
if i > 0 {
body.WriteByte(',')
}
_, _ = fmt.Fprintf(&body, `{"id":"model-%d"}`, i)
}
body.WriteString(`]}`)
return body.String()
}

View File

@@ -18,10 +18,10 @@ import (
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/test/bufconn"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
mgmtgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
"github.com/netbirdio/netbird/management/server/store"
"github.com/netbirdio/netbird/proxy/internal/middleware"
"github.com/netbirdio/netbird/proxy/internal/middleware/bodytap"

View File

@@ -1964,3 +1964,4 @@ func TestRouter_Serve_BacksOffOnTransientError(t *testing.T) {
type errSentinel string
func (e errSentinel) Error() string { return string(e) }

View File

@@ -271,6 +271,10 @@ func (c *testProxyController) SendServiceUpdateToCluster(_ context.Context, _ st
// noop
}
func (c *testProxyController) DiscoverModels(_ context.Context, _, _ string, _ *proto.ModelDiscoveryRequest) (*proto.ModelDiscoveryResult, error) {
return nil, nbproxy.ErrModelDiscoveryUnavailable
}
func (c *testProxyController) GetOIDCValidationConfig() nbproxy.OIDCValidationConfig {
return nbproxy.OIDCValidationConfig{}
}

View File

@@ -0,0 +1,349 @@
package proxy
import (
"context"
"errors"
"fmt"
"io"
"sync/atomic"
"testing"
"time"
log "github.com/sirupsen/logrus"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/grpc"
"github.com/netbirdio/netbird/proxy/internal/crowdsec"
"github.com/netbirdio/netbird/proxy/internal/modeldiscovery"
"github.com/netbirdio/netbird/shared/management/proto"
)
type stubModelDiscoverer struct {
started chan modeldiscovery.Request
release <-chan struct{}
result modeldiscovery.Result
err error
}
func (s *stubModelDiscoverer) Discover(ctx context.Context, request modeldiscovery.Request) (modeldiscovery.Result, error) {
if s.started != nil {
select {
case s.started <- request:
case <-ctx.Done():
return modeldiscovery.Result{}, ctx.Err()
}
}
if s.release != nil {
select {
case <-s.release:
case <-ctx.Done():
return modeldiscovery.Result{}, ctx.Err()
}
}
return s.result, s.err
}
type modelDiscoverySyncStream struct {
grpc.ClientStream
ctx context.Context
recv chan *proto.SyncMappingsResponse
sent chan *proto.SyncMappingsRequest
sendWait time.Duration
sending atomic.Int32
overlap atomic.Bool
}
func newModelDiscoverySyncStream(ctx context.Context) *modelDiscoverySyncStream {
return &modelDiscoverySyncStream{
ctx: ctx,
recv: make(chan *proto.SyncMappingsResponse, 16),
sent: make(chan *proto.SyncMappingsRequest, 16),
}
}
func (s *modelDiscoverySyncStream) Send(message *proto.SyncMappingsRequest) error {
if s.sending.Add(1) != 1 {
s.overlap.Store(true)
}
defer s.sending.Add(-1)
if s.sendWait > 0 {
time.Sleep(s.sendWait)
}
select {
case s.sent <- message:
return nil
case <-s.ctx.Done():
return s.ctx.Err()
}
}
func (s *modelDiscoverySyncStream) Recv() (*proto.SyncMappingsResponse, error) {
select {
case message, ok := <-s.recv:
if !ok {
return nil, io.EOF
}
return message, nil
case <-s.ctx.Done():
return nil, s.ctx.Err()
}
}
func (s *modelDiscoverySyncStream) Context() context.Context {
return s.ctx
}
func completeModelDiscoveryInitialSync(t *testing.T, stream *modelDiscoverySyncStream) {
t.Helper()
stream.recv <- &proto.SyncMappingsResponse{InitialSyncComplete: true}
select {
case sent := <-stream.sent:
require.NotNil(t, sent.GetAck())
case <-time.After(time.Second):
t.Fatal("initial snapshot was not acknowledged")
}
}
func TestProxyCapabilitiesAdvertiseModelDiscovery(t *testing.T) {
t.Parallel()
server := &Server{
crowdsecRegistry: crowdsec.NewRegistry("", "", log.New().WithField("test", true)),
}
capabilities := server.proxyCapabilities()
require.NotNil(t, capabilities.SupportsModelDiscovery)
assert.True(t, capabilities.GetSupportsModelDiscovery())
}
func TestExecuteModelDiscoveryMapsControlShape(t *testing.T) {
t.Parallel()
discoverer := &stubModelDiscoverer{
result: modeldiscovery.Result{
Source: modeldiscovery.SourceOpenAIV1Models,
Models: []modeldiscovery.Model{
{ID: "llama3.2:latest", Label: "Llama 3.2"},
},
},
}
request := &proto.ModelDiscoveryRequest{
RequestId: "request-1",
UpstreamUrl: "http://ollama.internal:11434",
AuthHeaderName: "Authorization",
AuthHeaderValue: "Bearer secret",
SkipTlsVerify: true,
OllamaFallback: true,
}
result := executeModelDiscovery(context.Background(), discoverer, request)
assert.Equal(t, "request-1", result.GetRequestId())
assert.Equal(t, modeldiscovery.SourceOpenAIV1Models, result.GetSource())
require.Len(t, result.GetModels(), 1)
assert.Equal(t, "llama3.2:latest", result.GetModels()[0].GetId())
assert.Equal(t, "Llama 3.2", result.GetModels()[0].GetLabel())
discoverer.started = make(chan modeldiscovery.Request, 1)
_ = executeModelDiscovery(context.Background(), discoverer, request)
received := <-discoverer.started
assert.Equal(t, request.GetUpstreamUrl(), received.UpstreamURL)
assert.Equal(t, request.GetAuthHeaderName(), received.AuthHeaderName)
assert.Equal(t, request.GetAuthHeaderValue(), received.AuthHeaderValue)
assert.True(t, received.SkipTLSVerify)
assert.True(t, received.AllowOllamaFallback)
}
func TestExecuteModelDiscoveryReturnsSanitizedError(t *testing.T) {
t.Parallel()
result := executeModelDiscovery(context.Background(), &stubModelDiscoverer{
err: errors.New("model discovery request failed"),
}, &proto.ModelDiscoveryRequest{RequestId: "request-error"})
assert.Equal(t, "request-error", result.GetRequestId())
assert.Equal(t, "model discovery request failed", result.GetError())
assert.Empty(t, result.GetModels())
assert.Empty(t, result.GetSource())
}
func TestHandleSyncMappingsStreamRunsDiscoveryOutOfBand(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
release := make(chan struct{})
started := make(chan modeldiscovery.Request, 1)
server := &Server{
Logger: log.New(),
routerReady: closedChan(),
modelDiscoverer: &stubModelDiscoverer{
started: started,
release: release,
result: modeldiscovery.Result{
Source: modeldiscovery.SourceOpenAIV1Models,
Models: []modeldiscovery.Model{{ID: "model-a", Label: "model-a"}},
},
},
}
stream := newModelDiscoverySyncStream(ctx)
stream.sendWait = 10 * time.Millisecond
done := make(chan error, 1)
initialSyncDone := false
go func() {
done <- server.handleSyncMappingsStream(ctx, stream, &initialSyncDone, time.Time{})
}()
completeModelDiscoveryInitialSync(t, stream)
stream.recv <- &proto.SyncMappingsResponse{
ModelDiscoveryRequest: &proto.ModelDiscoveryRequest{
RequestId: "request-1",
UpstreamUrl: "http://ollama.internal:11434",
},
}
select {
case <-started:
case <-time.After(time.Second):
t.Fatal("model discovery did not start")
}
// A normal mapping batch must still be acknowledged while the HTTP probe
// is in flight.
stream.recv <- &proto.SyncMappingsResponse{}
select {
case sent := <-stream.sent:
assert.NotNil(t, sent.GetAck())
assert.Nil(t, sent.GetModelDiscoveryResult())
case <-time.After(time.Second):
t.Fatal("mapping ack was blocked by model discovery")
}
close(release)
select {
case sent := <-stream.sent:
result := sent.GetModelDiscoveryResult()
require.NotNil(t, result)
assert.Equal(t, "request-1", result.GetRequestId())
assert.Equal(t, modeldiscovery.SourceOpenAIV1Models, result.GetSource())
assert.Nil(t, sent.GetAck())
case <-time.After(time.Second):
t.Fatal("model discovery result was not sent")
}
assert.False(t, stream.overlap.Load(), "acks and discovery results must use one serialized sender")
close(stream.recv)
require.NoError(t, <-done)
}
func TestHandleSyncMappingsStreamBoundsConcurrentDiscovery(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
release := make(chan struct{})
started := make(chan modeldiscovery.Request, 8)
server := &Server{
Logger: log.New(),
routerReady: closedChan(),
modelDiscoverer: &stubModelDiscoverer{
started: started,
release: release,
},
}
stream := newModelDiscoverySyncStream(ctx)
done := make(chan error, 1)
initialSyncDone := false
go func() {
done <- server.handleSyncMappingsStream(ctx, stream, &initialSyncDone, time.Time{})
}()
completeModelDiscoveryInitialSync(t, stream)
for i := range 5 {
stream.recv <- &proto.SyncMappingsResponse{
ModelDiscoveryRequest: &proto.ModelDiscoveryRequest{
RequestId: fmt.Sprintf("request-%d", i),
UpstreamUrl: "http://ollama.internal:11434",
},
}
}
for range 4 {
select {
case <-started:
case <-time.After(time.Second):
t.Fatal("expected four concurrent model discoveries")
}
}
select {
case sent := <-stream.sent:
result := sent.GetModelDiscoveryResult()
require.NotNil(t, result)
assert.Equal(t, "model discovery is busy", result.GetError())
case <-time.After(time.Second):
t.Fatal("fifth discovery did not fail fast")
}
close(release)
for range 4 {
select {
case sent := <-stream.sent:
require.NotNil(t, sent.GetModelDiscoveryResult())
case <-time.After(time.Second):
t.Fatal("in-flight model discovery did not complete")
}
}
close(stream.recv)
require.NoError(t, <-done)
}
func TestHandleSyncMappingsStreamRejectsMixedDiscoveryMessage(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
server := &Server{
Logger: log.New(),
routerReady: closedChan(),
modelDiscoverer: &stubModelDiscoverer{},
}
stream := newModelDiscoverySyncStream(ctx)
stream.recv <- &proto.SyncMappingsResponse{
Mapping: []*proto.ProxyMapping{{Id: "mapping-1"}},
InitialSyncComplete: true,
ModelDiscoveryRequest: &proto.ModelDiscoveryRequest{
RequestId: "request-1",
},
}
close(stream.recv)
initialSyncDone := true
err := server.handleSyncMappingsStream(ctx, stream, &initialSyncDone, time.Time{})
require.EqualError(t, err, "model discovery message must not include mapping data or set initial_sync_complete")
}
func TestHandleSyncMappingsStreamRejectsDiscoveryBeforeInitialSnapshot(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
server := &Server{
Logger: log.New(),
routerReady: closedChan(),
modelDiscoverer: &stubModelDiscoverer{},
}
stream := newModelDiscoverySyncStream(ctx)
stream.recv <- &proto.SyncMappingsResponse{
ModelDiscoveryRequest: &proto.ModelDiscoveryRequest{
RequestId: "request-1",
},
}
close(stream.recv)
initialSyncDone := true
err := server.handleSyncMappingsStream(ctx, stream, &initialSyncDone, time.Time{})
require.EqualError(t, err, "model discovery request received before initial sync completed")
}

View File

@@ -57,6 +57,7 @@ import (
proxymetrics "github.com/netbirdio/netbird/proxy/internal/metrics"
"github.com/netbirdio/netbird/proxy/internal/middleware"
mwbuiltin "github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
"github.com/netbirdio/netbird/proxy/internal/modeldiscovery"
"github.com/netbirdio/netbird/proxy/internal/netutil"
"github.com/netbirdio/netbird/proxy/internal/proxy"
"github.com/netbirdio/netbird/proxy/internal/restrict"
@@ -78,6 +79,10 @@ type portRouter struct {
cancel context.CancelFunc
}
type providerModelDiscoverer interface {
Discover(context.Context, modeldiscovery.Request) (modeldiscovery.Result, error)
}
type Server struct {
ctx context.Context
mgmtClient proto.ProxyServiceClient
@@ -100,16 +105,20 @@ type Server struct {
// middlewareRegistry is the source of registered middleware factories.
// Concrete middlewares register themselves through init().
middlewareRegistry *middleware.Registry
mainRouter *nbtcp.Router
mainPort uint16
udpMu sync.Mutex
udpRelays map[types.ServiceID]*udprelay.Relay
udpRelayWg sync.WaitGroup
portMu sync.RWMutex
portRouters map[uint16]*portRouter
svcPorts map[types.ServiceID][]uint16
lastMappings map[types.ServiceID]*proto.ProxyMapping
portRouterWg sync.WaitGroup
// modelDiscoverer executes explicit provider model-list probes on the
// proxy's host network. Lazily constructed for normal servers; injectable
// in focused control-stream tests.
modelDiscoverer providerModelDiscoverer
mainRouter *nbtcp.Router
mainPort uint16
udpMu sync.Mutex
udpRelays map[types.ServiceID]*udprelay.Relay
udpRelayWg sync.WaitGroup
portMu sync.RWMutex
portRouters map[uint16]*portRouter
svcPorts map[types.ServiceID][]uint16
lastMappings map[types.ServiceID]*proto.ProxyMapping
portRouterWg sync.WaitGroup
// hijackTracker tracks hijacked connections (e.g. WebSocket upgrades)
// so they can be closed during graceful shutdown, since http.Server.Shutdown
@@ -1277,12 +1286,16 @@ func (s *Server) proxyCapabilities() *proto.ProxyCapabilities {
privateCapability := s.Private
// Always true: this build enforces ProxyMapping.private via the auth middleware.
supportsPrivateService := true
// Model discovery is handled only on SyncMappings. Management also gates
// dispatch on the live connection using this capability.
supportsModelDiscovery := true
return &proto.ProxyCapabilities{
SupportsCustomPorts: &s.SupportsCustomPorts,
RequireSubdomain: &s.RequireSubdomain,
SupportsCrowdsec: &supportsCrowdSec,
Private: &privateCapability,
SupportsPrivateService: &supportsPrivateService,
SupportsModelDiscovery: &supportsModelDiscovery,
}
}
@@ -1349,7 +1362,8 @@ func isSyncUnimplemented(err error) bool {
// handleSyncMappingsStream consumes batches from a bidirectional SyncMappings
// stream, sending an ack after each batch is fully processed. Management waits
// for the ack before sending the next batch, providing application-level
// back-pressure.
// back-pressure. Model discovery commands are out-of-band: they run with
// bounded concurrency and return a correlated result instead of an ack.
func (s *Server) handleSyncMappingsStream(ctx context.Context, stream proto.ProxyService_SyncMappingsClient, initialSyncDone *bool, connectTime time.Time) error {
select {
case <-s.routerReady:
@@ -1358,6 +1372,30 @@ func (s *Server) handleSyncMappingsStream(ctx context.Context, stream proto.Prox
}
tracker := s.newSnapshotTracker(initialSyncDone, connectTime)
initialSnapshotComplete := false
discoverer := s.modelDiscoverer
if discoverer == nil {
discoverer = modeldiscovery.New(s.Logger)
}
const maxConcurrentDiscoveries = 4
discoverySlots := make(chan struct{}, maxConcurrentDiscoveries)
discoveryCtx, cancelDiscoveries := context.WithCancel(ctx)
var discoveryWG sync.WaitGroup
var sendMu sync.Mutex
defer func() {
cancelDiscoveries()
discoveryWG.Wait()
}()
// gRPC permits one concurrent sender and one concurrent receiver, but not
// multiple senders. Mapping acks and asynchronous discovery results share
// this serialized send path.
send := func(message *proto.SyncMappingsRequest) error {
sendMu.Lock()
defer sendMu.Unlock()
return stream.Send(message)
}
for {
select {
@@ -1372,15 +1410,62 @@ func (s *Server) handleSyncMappingsStream(ctx context.Context, stream proto.Prox
return fmt.Errorf("receive msg: %w", err)
}
if discovery := msg.GetModelDiscoveryRequest(); discovery != nil {
if len(msg.GetMapping()) != 0 || msg.GetInitialSyncComplete() {
return errors.New("model discovery message must not include mapping data or set initial_sync_complete")
}
if !initialSnapshotComplete {
return errors.New("model discovery request received before initial sync completed")
}
select {
case discoverySlots <- struct{}{}:
discoveryWG.Add(1)
go func() {
defer discoveryWG.Done()
defer func() { <-discoverySlots }()
result := executeModelDiscovery(discoveryCtx, discoverer, discovery)
if discoveryCtx.Err() != nil {
return
}
if err := send(&proto.SyncMappingsRequest{
Msg: &proto.SyncMappingsRequest_ModelDiscoveryResult{
ModelDiscoveryResult: result,
},
}); err != nil {
s.Logger.WithError(err).Debug("failed to send model discovery result")
}
}()
default:
result := &proto.ModelDiscoveryResult{
RequestId: discovery.GetRequestId(),
Error: "model discovery is busy",
}
if err := send(&proto.SyncMappingsRequest{
Msg: &proto.SyncMappingsRequest_ModelDiscoveryResult{
ModelDiscoveryResult: result,
},
}); err != nil {
return fmt.Errorf("send model discovery busy result: %w", err)
}
}
continue
}
batchStart := time.Now()
s.Logger.Debug("Received mapping update, starting processing")
if err := s.processMappingsGuarded(ctx, msg.GetMapping()); err != nil {
return err
}
s.Logger.Debug("Processing mapping update completed")
tracker.recordBatch(ctx, s, msg.GetMapping(), msg.GetInitialSyncComplete(), batchStart)
syncComplete := msg.GetInitialSyncComplete()
tracker.recordBatch(ctx, s, msg.GetMapping(), syncComplete, batchStart)
if syncComplete {
initialSnapshotComplete = true
}
if err := stream.Send(&proto.SyncMappingsRequest{
if err := send(&proto.SyncMappingsRequest{
Msg: &proto.SyncMappingsRequest_Ack{Ack: &proto.SyncMappingsAck{}},
}); err != nil {
return fmt.Errorf("send ack: %w", err)
@@ -1389,6 +1474,31 @@ func (s *Server) handleSyncMappingsStream(ctx context.Context, stream proto.Prox
}
}
func executeModelDiscovery(ctx context.Context, discoverer providerModelDiscoverer, request *proto.ModelDiscoveryRequest) *proto.ModelDiscoveryResult {
result := &proto.ModelDiscoveryResult{RequestId: request.GetRequestId()}
discovered, err := discoverer.Discover(ctx, modeldiscovery.Request{
UpstreamURL: request.GetUpstreamUrl(),
AuthHeaderName: request.GetAuthHeaderName(),
AuthHeaderValue: request.GetAuthHeaderValue(),
SkipTLSVerify: request.GetSkipTlsVerify(),
AllowOllamaFallback: request.GetOllamaFallback(),
})
if err != nil {
result.Error = err.Error()
return result
}
result.Source = discovered.Source
result.Models = make([]*proto.ModelDiscoveryModel, 0, len(discovered.Models))
for _, model := range discovered.Models {
result.Models = append(result.Models, &proto.ModelDiscoveryModel{
Id: model.ID,
Label: model.Label,
})
}
return result
}
// snapshotTracker accumulates service IDs during the initial snapshot and
// finalises sync state when the complete flag arrives. Used by both
// handleMappingStream and handleSyncMappingsStream so metric emission and

View File

@@ -5138,6 +5138,9 @@ components:
description: Models exposed through this endpoint, with the operator's per-1k input/output prices. Empty means all catalog models are allowed at catalog prices.
items:
$ref: '#/components/schemas/AgentNetworkProviderModel'
has_api_key:
type: boolean
description: Whether an upstream API key is currently stored. The key itself is never returned.
extra_values:
type: object
description: |
@@ -5186,6 +5189,7 @@ components:
- name
- upstream_url
- models
- has_api_key
- enabled
- skip_tls_verification
- metadata_disabled
@@ -5212,7 +5216,8 @@ components:
example: "eu.proxy.netbird.io"
api_key:
type: string
description: Upstream provider API key. Sealed at rest on the management server and never returned in responses. Required on create; optional on update (omit to keep the existing key).
description: |
Upstream provider API key. Sealed at rest on the management server and never returned in responses. Whether a key is accepted or required is declared by the selected catalog provider's `auth_mode`. On update, omit this field to preserve the existing key or send an empty string to clear an optional key.
example: "sk-..."
models:
type: array
@@ -5275,6 +5280,51 @@ components:
- id
- input_per_1k
- output_per_1k
AgentNetworkDiscoveredModel:
type: object
description: A model reported by the provider's persisted upstream endpoint.
properties:
id:
type: string
description: Exact model identifier accepted by the upstream endpoint.
example: "llama3.2:latest"
label:
type: string
description: Human-friendly label reported by the endpoint. Falls back to the model identifier.
example: "llama3.2:latest"
required:
- id
- label
AgentNetworkModelDiscoveryResponse:
type: object
description: |
Result of an explicit model-discovery probe executed by a connected
proxy in the account's selected Agent Network cluster. Discovered
models are not persisted until the provider is updated.
properties:
models:
type: array
description: Normalized, deduplicated models reported by the upstream.
items:
$ref: '#/components/schemas/AgentNetworkDiscoveredModel'
source:
type: string
description: Upstream API shape used to obtain the result.
enum: [openai_v1_models, ollama_api_tags]
example: "openai_v1_models"
proxy_cluster:
type: string
description: Proxy cluster from which the endpoint was queried.
example: "eu.proxy.netbird.io"
request_id:
type: string
description: Correlation identifier for the management-to-proxy probe.
example: "b7ef7da0-78cc-4583-8c4f-55f2ce72015f"
required:
- models
- source
- proxy_cluster
- request_id
AgentNetworkCatalogModel:
type: object
properties:
@@ -5325,6 +5375,11 @@ components:
type: string
description: Default upstream host suggested when adding a provider of this type.
example: "api.openai.com"
auth_mode:
type: string
description: Whether this provider requires, optionally accepts, or does not support an upstream API key.
enum: [required, optional, none]
example: "required"
auth_header_template:
type: string
description: Template the proxy uses to inject the API key (the literal string ${API_KEY} is replaced at request time).
@@ -5346,6 +5401,10 @@ components:
"custom" — generic OpenAI-compatible self-hosted endpoint catch-all.
enum: [provider, gateway, custom]
example: "provider"
supports_model_discovery:
type: boolean
description: Whether models can be loaded from a persisted endpoint through the selected proxy cluster.
example: false
extra_headers:
type: array
description: |
@@ -5364,10 +5423,12 @@ components:
- name
- description
- default_host
- auth_mode
- auth_header_template
- default_content_type
- brand_color
- kind
- supports_model_discovery
- models
AgentNetworkCatalogIdentityInjection:
type: object
@@ -14028,6 +14089,48 @@ paths:
"$ref": "#/components/responses/not_found"
'500':
"$ref": "#/components/responses/internal_error"
/api/agent-network/providers/{providerId}/discover-models:
post:
summary: Discover models from an Agent Network provider
description: |
Executes a bounded, read-only model-list probe from a connected proxy
in the account's selected Agent Network cluster. The probe uses the
provider's persisted endpoint, TLS setting, and stored credential.
Returned models are not persisted by this operation.
tags: [ Agent Network ]
security:
- BearerAuth: [ ]
- TokenAuth: [ ]
parameters:
- in: path
name: providerId
required: true
schema:
type: string
description: The unique identifier of a persisted Agent Network provider
responses:
'200':
description: Models discovered successfully
content:
application/json:
schema:
$ref: '#/components/schemas/AgentNetworkModelDiscoveryResponse'
'400':
"$ref": "#/components/responses/bad_request"
'401':
"$ref": "#/components/responses/requires_authentication"
'403':
"$ref": "#/components/responses/forbidden"
'404':
"$ref": "#/components/responses/not_found"
'412':
description: The provider cannot be queried or no capable proxy is connected
content: { }
'429':
description: A discovery probe is already running or was requested too recently
content: { }
'500':
"$ref": "#/components/responses/internal_error"
/api/agent-network/policies:
get:
summary: List all Agent Network Policies

View File

@@ -38,6 +38,27 @@ func (e AccessRestrictionsCrowdsecMode) Valid() bool {
}
}
// Defines values for AgentNetworkCatalogProviderAuthMode.
const (
AgentNetworkCatalogProviderAuthModeNone AgentNetworkCatalogProviderAuthMode = "none"
AgentNetworkCatalogProviderAuthModeOptional AgentNetworkCatalogProviderAuthMode = "optional"
AgentNetworkCatalogProviderAuthModeRequired AgentNetworkCatalogProviderAuthMode = "required"
)
// Valid indicates whether the value is a known member of the AgentNetworkCatalogProviderAuthMode enum.
func (e AgentNetworkCatalogProviderAuthMode) Valid() bool {
switch e {
case AgentNetworkCatalogProviderAuthModeNone:
return true
case AgentNetworkCatalogProviderAuthModeOptional:
return true
case AgentNetworkCatalogProviderAuthModeRequired:
return true
default:
return false
}
}
// Defines values for AgentNetworkCatalogProviderKind.
const (
AgentNetworkCatalogProviderKindCustom AgentNetworkCatalogProviderKind = "custom"
@@ -77,6 +98,24 @@ func (e AgentNetworkConsumptionDimensionKind) Valid() bool {
}
}
// Defines values for AgentNetworkModelDiscoveryResponseSource.
const (
AgentNetworkModelDiscoveryResponseSourceOllamaApiTags AgentNetworkModelDiscoveryResponseSource = "ollama_api_tags"
AgentNetworkModelDiscoveryResponseSourceOpenaiV1Models AgentNetworkModelDiscoveryResponseSource = "openai_v1_models"
)
// Valid indicates whether the value is a known member of the AgentNetworkModelDiscoveryResponseSource enum.
func (e AgentNetworkModelDiscoveryResponseSource) Valid() bool {
switch e {
case AgentNetworkModelDiscoveryResponseSourceOllamaApiTags:
return true
case AgentNetworkModelDiscoveryResponseSourceOpenaiV1Models:
return true
default:
return false
}
}
// Defines values for CreateAzureIntegrationRequestHost.
const (
CreateAzureIntegrationRequestHostMicrosoftCom CreateAzureIntegrationRequestHost = "microsoft.com"
@@ -2038,6 +2077,9 @@ type AgentNetworkCatalogProvider struct {
// AuthHeaderTemplate Template the proxy uses to inject the API key (the literal string ${API_KEY} is replaced at request time).
AuthHeaderTemplate string `json:"auth_header_template"`
// AuthMode Whether this provider requires, optionally accepts, or does not support an upstream API key.
AuthMode AgentNetworkCatalogProviderAuthMode `json:"auth_mode"`
// BrandColor Hex brand color used to render the provider badge in the dashboard.
BrandColor string `json:"brand_color"`
@@ -2070,8 +2112,14 @@ type AgentNetworkCatalogProvider struct {
// Name Display name for the provider.
Name string `json:"name"`
// SupportsModelDiscovery Whether models can be loaded from a persisted endpoint through the selected proxy cluster.
SupportsModelDiscovery bool `json:"supports_model_discovery"`
}
// AgentNetworkCatalogProviderAuthMode Whether this provider requires, optionally accepts, or does not support an upstream API key.
type AgentNetworkCatalogProviderAuthMode string
// AgentNetworkCatalogProviderKind Presentation grouping for the provider Select on the dashboard.
// "provider" — first-party vendor API (OpenAI, Anthropic, …); the upstream is the model itself.
// "gateway" — routing/aggregation layer in front of multiple providers (LiteLLM, Portkey, …); typically pairs with NetBird identity stamping.
@@ -2108,6 +2156,15 @@ type AgentNetworkConsumption struct {
// AgentNetworkConsumptionDimensionKind Whether this row counts a single end user or a single source group across every member.
type AgentNetworkConsumptionDimensionKind string
// AgentNetworkDiscoveredModel A model reported by the provider's persisted upstream endpoint.
type AgentNetworkDiscoveredModel struct {
// Id Exact model identifier accepted by the upstream endpoint.
Id string `json:"id"`
// Label Human-friendly label reported by the endpoint. Falls back to the model identifier.
Label string `json:"label"`
}
// AgentNetworkGuardrail defines model for AgentNetworkGuardrail.
type AgentNetworkGuardrail struct {
// Checks Guardrail check parameters. Each entry has an `enabled` flag plus per-check configuration; disabled entries are inert.
@@ -2155,6 +2212,26 @@ type AgentNetworkGuardrailRequest struct {
Name string `json:"name"`
}
// AgentNetworkModelDiscoveryResponse Result of an explicit model-discovery probe executed by a connected
// proxy in the account's selected Agent Network cluster. Discovered
// models are not persisted until the provider is updated.
type AgentNetworkModelDiscoveryResponse struct {
// Models Normalized, deduplicated models reported by the upstream.
Models []AgentNetworkDiscoveredModel `json:"models"`
// ProxyCluster Proxy cluster from which the endpoint was queried.
ProxyCluster string `json:"proxy_cluster"`
// RequestId Correlation identifier for the management-to-proxy probe.
RequestId string `json:"request_id"`
// Source Upstream API shape used to obtain the result.
Source AgentNetworkModelDiscoveryResponseSource `json:"source"`
}
// AgentNetworkModelDiscoveryResponseSource Upstream API shape used to obtain the result.
type AgentNetworkModelDiscoveryResponseSource string
// AgentNetworkPolicy defines model for AgentNetworkPolicy.
type AgentNetworkPolicy struct {
// CreatedAt Timestamp when the policy was created.
@@ -2260,6 +2337,9 @@ type AgentNetworkProvider struct {
// ExtraValues Operator-typed values for catalog-declared extra headers. Keys are wire header names (e.g. `x-portkey-config`); values are the strings the proxy stamps on every upstream request to this provider. Catalog (AgentNetworkCatalogProvider.extra_headers) declares which keys are accepted; values not declared by the catalog are ignored at synth time. Empty / missing values mean no header stamped.
ExtraValues *map[string]string `json:"extra_values,omitempty"`
// HasApiKey Whether an upstream API key is currently stored. The key itself is never returned.
HasApiKey bool `json:"has_api_key"`
// Id Provider ID
Id string `json:"id"`
@@ -2305,7 +2385,7 @@ type AgentNetworkProviderModel struct {
// AgentNetworkProviderRequest defines model for AgentNetworkProviderRequest.
type AgentNetworkProviderRequest struct {
// ApiKey Upstream provider API key. Sealed at rest on the management server and never returned in responses. Required on create; optional on update (omit to keep the existing key).
// ApiKey Upstream provider API key. Sealed at rest on the management server and never returned in responses. Whether a key is accepted or required is declared by the selected catalog provider's `auth_mode`. On update, omit this field to preserve the existing key or send an empty string to clear an optional key.
ApiKey *string `json:"api_key,omitempty"`
// BootstrapCluster Proxy cluster used to bootstrap the per-account agent-network endpoint when the first provider is created. Ignored on subsequent creates and on updates because the cluster is pinned on the account-level Settings row.

File diff suppressed because it is too large Load Diff

View File

@@ -73,6 +73,10 @@ message ProxyCapabilities {
optional bool private = 4;
// Whether the proxy enforces ProxyMapping.private (fails closed on ValidateTunnelPeer failure). Management MUST NOT stream private mappings to proxies that don't claim this.
optional bool supports_private_service = 5;
// Whether the proxy can execute correlated Agent Network model-discovery
// requests delivered over SyncMappings. Management must not send discovery
// requests to proxies that omit or disable this capability.
optional bool supports_model_discovery = 6;
}
// GetMappingUpdateRequest is sent to initialise a mapping stream.
@@ -420,6 +424,8 @@ message SyncMappingsRequest {
oneof msg {
SyncMappingsInit init = 1;
SyncMappingsAck ack = 2;
// Correlated response to a ModelDiscoveryRequest received from management.
ModelDiscoveryResult model_discovery_result = 3;
}
}
@@ -443,6 +449,9 @@ message SyncMappingsResponse {
repeated ProxyMapping mapping = 1;
// initial_sync_complete is set on the last message of the initial snapshot.
bool initial_sync_complete = 2;
// An out-of-band model-discovery operation. This is sent only after the
// initial snapshot and only to proxies advertising supports_model_discovery.
ModelDiscoveryRequest model_discovery_request = 3;
}
// CheckLLMPolicyLimitsRequest carries the resolved caller identity and the
@@ -501,3 +510,33 @@ message RecordLLMUsageRequest {
message RecordLLMUsageResponse {
}
// ModelDiscoveryRequest asks a proxy to list models from a persisted provider
// endpoint using the proxy host's network path. request_id is assigned by
// management and echoed verbatim in ModelDiscoveryResult.
message ModelDiscoveryRequest {
string request_id = 1;
string upstream_url = 2;
string auth_header_name = 3;
string auth_header_value = 4;
bool skip_tls_verify = 5;
// When true, the proxy may fall back from the OpenAI-compatible
// GET /v1/models endpoint to Ollama's native GET /api/tags endpoint.
bool ollama_fallback = 6;
}
// ModelDiscoveryModel is the normalized model shape returned to management.
// Arbitrary upstream fields never cross the proxy control channel.
message ModelDiscoveryModel {
string id = 1;
string label = 2;
}
// ModelDiscoveryResult completes one ModelDiscoveryRequest. error is empty on
// success and contains only a sanitized diagnostic on failure.
message ModelDiscoveryResult {
string request_id = 1;
repeated ModelDiscoveryModel models = 2;
// Stable source label, currently openai_v1_models or ollama_api_tags.
string source = 3;
string error = 4;
}