mirror of
https://github.com/netbirdio/netbird.git
synced 2026-07-28 11:21:29 +02:00
Compare commits
5 Commits
mlsmaycon-
...
feat/agent
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
fea131101d | ||
|
|
0463b30152 | ||
|
|
d3909e4faf | ||
|
|
a92cdb7dcd | ||
|
|
564abb660c |
@@ -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))
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
124
management/internals/modules/agentnetwork/provider_auth_test.go
Normal file
124
management/internals/modules/agentnetwork/provider_auth_test.go
Normal 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")
|
||||
})
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
|
||||
404
management/internals/shared/grpc/model_discovery_test.go
Normal file
404
management/internals/shared/grpc/model_discovery_test.go
Normal 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())
|
||||
}
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
@@ -29,6 +30,7 @@ import (
|
||||
|
||||
"github.com/netbirdio/netbird/shared/management/domain"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/peers"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
|
||||
@@ -36,7 +38,6 @@ import (
|
||||
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/sessionkey"
|
||||
"github.com/netbirdio/netbird/management/server/idp"
|
||||
"github.com/netbirdio/netbird/management/server/peer"
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
|
||||
"github.com/netbirdio/netbird/management/server/types"
|
||||
"github.com/netbirdio/netbird/management/server/users"
|
||||
proxyauth "github.com/netbirdio/netbird/proxy/auth"
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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))
|
||||
|
||||
337
proxy/internal/modeldiscovery/discovery.go
Normal file
337
proxy/internal/modeldiscovery/discovery.go
Normal 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
|
||||
}
|
||||
315
proxy/internal/modeldiscovery/discovery_test.go
Normal file
315
proxy/internal/modeldiscovery/discovery_test.go
Normal 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()
|
||||
}
|
||||
@@ -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{}
|
||||
}
|
||||
|
||||
349
proxy/model_discovery_sync_test.go
Normal file
349
proxy/model_discovery_sync_test.go
Normal 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")
|
||||
}
|
||||
136
proxy/server.go
136
proxy/server.go
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user