diff --git a/management/internals/modules/agentnetwork/synthesizer.go b/management/internals/modules/agentnetwork/synthesizer.go index c7fab2241..e7e6e730b 100644 --- a/management/internals/modules/agentnetwork/synthesizer.go +++ b/management/internals/modules/agentnetwork/synthesizer.go @@ -935,6 +935,12 @@ func buildAccountService( middlewares []rpservice.MiddlewareConfig, sessionPriv, sessionPub string, ) *rpservice.Service { + // The proxy that serves this tenant — a dedicated proxy when one has been + // assigned, else the shared cluster. This is the value mesh-DNS peer + // selection and the connect-snapshot filter both join on. + servingProxy := settings.ServingProxy() + // The shared cluster address remains the placeholder target's ID; only the + // advertised proxy address follows ServingProxy(). cluster := settings.Cluster domain := settings.Endpoint() serviceID := SynthesizedServiceIDPrefix + accountID @@ -944,7 +950,7 @@ func buildAccountService( AccountID: accountID, Name: "agent-network-" + accountID, Domain: domain, - ProxyCluster: cluster, + ProxyCluster: servingProxy, DNSZone: settings.Zone, // empty for legacy rows → unchanged behavior Mode: rpservice.ModeHTTP, Enabled: true, diff --git a/management/internals/modules/agentnetwork/synthesizer_test.go b/management/internals/modules/agentnetwork/synthesizer_test.go index 7b14f8209..8a3a65d30 100644 --- a/management/internals/modules/agentnetwork/synthesizer_test.go +++ b/management/internals/modules/agentnetwork/synthesizer_test.go @@ -1246,3 +1246,99 @@ 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") } + +// TestBuildAccountService_ProxyClusterFollowsServingProxyAddress — the whole +// point of the column: the synthesized service must advertise the private +// proxy's address, because that value is what mesh-DNS peer selection and the +// connect-snapshot filter both join on. TargetId must NOT move with it — it +// names the noop placeholder target and is out of scope (see the note above). +func TestBuildAccountService_ProxyClusterFollowsServingProxyAddress(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + settings := &types.Settings{ + AccountID: testAccountID, + Cluster: testCluster, + Zone: "gateway.netbird.ai", + Subdomain: "brave-otter", + ServingProxyAddress: "brave-otter.gateway.netbird.ai", + } + provider := newSynthTestProvider() + policy := newSynthTestPolicy(provider.ID, "grp-eng", "") + + expectSynthBaseInputs(mockStore, ctx, settings, + []*types.Provider{provider}, + []*types.Policy{policy}, + []*types.Guardrail{}) + + services, err := SynthesizeServices(ctx, mockStore, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1) + + svc := services[0] + assert.Equal(t, "brave-otter.gateway.netbird.ai", svc.ProxyCluster, + "ProxyCluster must advertise the private proxy's address once ServingProxyAddress is set") + require.Len(t, svc.Targets, 1) + assert.Equal(t, testCluster, svc.Targets[0].TargetId, + "TargetId is the noop placeholder target and must stay pinned to the shared cluster, not the serving proxy") +} + +// TestSynthesizeServicesForCluster_ExcludesPrivatelyServedTenant — a tenant +// moved to a private proxy must drop out of the SHARED proxy's connect +// snapshot, or both proxies would serve it. The existing +// `svc.ProxyCluster == clusterAddr` filter does this for free once ProxyCluster +// is the tenant hostname; this test proves the handoff rather than assuming it. +func TestSynthesizeServicesForCluster_ExcludesPrivatelyServedTenant(t *testing.T) { + ctx := context.Background() + + provider := newSynthTestProvider() + policy := newSynthTestPolicy(provider.ID, "grp-eng", "") + + privatelyServed := &types.Settings{ + AccountID: testAccountID, + Cluster: testCluster, + Subdomain: testSubdomain, + ServingProxyAddress: "brave-otter.gateway.netbird.ai", + } + + t.Run("privately served tenant is excluded from the shared cluster snapshot", func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + mockStore.EXPECT(). + GetAgentNetworkSettingsByCluster(ctx, store.LockingStrengthNone, testCluster). + Return([]*types.Settings{privatelyServed}, nil) + expectSynthBaseInputs(mockStore, ctx, privatelyServed, + []*types.Provider{provider}, []*types.Policy{policy}, []*types.Guardrail{}) + + services, err := SynthesizeServicesForCluster(ctx, mockStore, testCluster) + require.NoError(t, err) + assert.Empty(t, services, "a tenant served by a private proxy must not appear in the shared cluster's snapshot") + }) + + t.Run("clearing ServingProxyAddress makes the tenant reappear", func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockStore := store.NewMockStore(ctrl) + + sharedAgain := &types.Settings{ + AccountID: testAccountID, + Cluster: testCluster, + Subdomain: testSubdomain, + } + + mockStore.EXPECT(). + GetAgentNetworkSettingsByCluster(ctx, store.LockingStrengthNone, testCluster). + Return([]*types.Settings{sharedAgain}, nil) + expectSynthBaseInputs(mockStore, ctx, sharedAgain, + []*types.Provider{provider}, []*types.Policy{policy}, []*types.Guardrail{}) + + services, err := SynthesizeServicesForCluster(ctx, mockStore, testCluster) + require.NoError(t, err) + require.Len(t, services, 1, "clearing ServingProxyAddress must return the tenant to the shared cluster's snapshot") + assert.Equal(t, testCluster, services[0].ProxyCluster) + }) +} diff --git a/management/internals/modules/agentnetwork/types/settings.go b/management/internals/modules/agentnetwork/types/settings.go index ba707bb08..661ed1051 100644 --- a/management/internals/modules/agentnetwork/types/settings.go +++ b/management/internals/modules/agentnetwork/types/settings.go @@ -11,9 +11,11 @@ import ( // the long-term aggregate and are retained independently. const DefaultAccessLogRetentionDays = 30 -// Settings is the per-account agent-network configuration row. One -// row per account. Cluster + Subdomain are immutable once written and -// produce the public endpoint agents call (`.`). +// Settings is the per-account agent-network configuration row. One row per +// account. The public endpoint agents call is `.` when a +// zone is set, else `.`. Cluster, Subdomain and Zone are +// immutable once written; ServingProxyAddress is the one mutable column, +// naming which proxy currently serves the account. type Settings struct { AccountID string `gorm:"primaryKey"` Cluster string @@ -27,6 +29,20 @@ type Settings struct { // configures no zone keep that behaviour unchanged. Zone string + // ServingProxyAddress is the address of the proxy currently serving this + // account's gateway. Empty means the account is served by the shared proxy + // at Cluster; set means a dedicated proxy serves it, and the value is that + // proxy's address — for a per-account proxy, the account's own gateway + // hostname. + // + // This is the only mutable column on this row. Cluster, Subdomain and Zone + // are fixed once written, but moving an account onto a dedicated proxy — and + // moving it back — is exactly one write here. Nothing in this repository + // writes it: it is set by whatever external process assigns dedicated + // proxies, and its zero value preserves existing behaviour for every current + // row and every deployment that assigns none. + ServingProxyAddress string + // Account-level collection controls sourced by the synthesizer. // EnableLogCollection gates the per-request access-log trail and defaults // ON for new accounts. EnablePromptCollection is the master gate for @@ -64,6 +80,17 @@ func (s *Settings) Endpoint() string { return s.Subdomain + "." + s.Cluster } +// ServingProxy returns the address of the proxy that serves this account's +// gateway: the dedicated proxy when one has been assigned, otherwise the shared +// cluster. This is the value the synthesized service advertises as +// ProxyCluster, which is what mesh-DNS peer selection joins on. +func (s *Settings) ServingProxy() string { + if s.ServingProxyAddress != "" { + return s.ServingProxyAddress + } + return s.Cluster +} + // ToAPIResponse renders the settings as the API representation. func (s *Settings) ToAPIResponse() *api.AgentNetworkSettings { created := s.CreatedAt diff --git a/management/internals/modules/agentnetwork/types/settings_test.go b/management/internals/modules/agentnetwork/types/settings_test.go index 8837cb7d5..41f7fd55b 100644 --- a/management/internals/modules/agentnetwork/types/settings_test.go +++ b/management/internals/modules/agentnetwork/types/settings_test.go @@ -29,3 +29,17 @@ func TestToAPIResponse_ExposesZoneAndDerivedEndpoint(t *testing.T) { resp := s.ToAPIResponse() assert.Equal(t, "brave-otter.gateway.netbird.ai", resp.Endpoint) } + +// TestServingProxy_PrefersColumnOverCluster — a provisioned tenant is served by +// its own proxy, whose address is its hostname, not the shared cluster. +func TestServingProxy_PrefersColumnOverCluster(t *testing.T) { + s := &Settings{Cluster: "eu.proxy.netbird.io", ServingProxyAddress: "brave-otter.gateway.netbird.ai"} + assert.Equal(t, "brave-otter.gateway.netbird.ai", s.ServingProxy()) +} + +// TestServingProxy_FallsBackToCluster is the compatibility guarantee: every +// existing row, and every self-hosted deployment, is served by the shared proxy. +func TestServingProxy_FallsBackToCluster(t *testing.T) { + s := &Settings{Cluster: "eu.proxy.netbird.io"} + assert.Equal(t, "eu.proxy.netbird.io", s.ServingProxy()) +} diff --git a/management/server/store/sql_store_agentnetwork.go b/management/server/store/sql_store_agentnetwork.go index 883923156..685071125 100644 --- a/management/server/store/sql_store_agentnetwork.go +++ b/management/server/store/sql_store_agentnetwork.go @@ -385,6 +385,24 @@ func (s *SqlStore) CreateAgentNetworkSettings(ctx context.Context, settings *age return nil } +// SetAgentNetworkServingProxyAddress points the account's gateway at a specific +// serving proxy, or clears it (address == "") to return the account to the +// shared proxy. Scoped to the one column on purpose: this runs concurrently +// with unrelated settings updates, and a full-row upsert would clobber them. +func (s *SqlStore) SetAgentNetworkServingProxyAddress(ctx context.Context, accountID, address string) error { + result := s.db.Model(&agentNetworkTypes.Settings{}). + Where("account_id = ?", accountID). + Update("serving_proxy_address", address) + if result.Error != nil { + log.WithContext(ctx).Errorf("failed to set agent network serving proxy address: %v", result.Error) + return status.Errorf(status.Internal, "failed to set agent network serving proxy address") + } + if result.RowsAffected == 0 { + return status.Errorf(status.NotFound, "agent network settings for account %s not found", accountID) + } + return nil +} + // IncrementAgentNetworkConsumption atomically upserts the consumption // row keyed on (account, dim_kind, dim_id, window_seconds, window_start) // and adds the supplied deltas. Concurrent calls from multiple proxy diff --git a/management/server/store/store.go b/management/server/store/store.go index 3aab631e4..dc685d16e 100644 --- a/management/server/store/store.go +++ b/management/server/store/store.go @@ -364,6 +364,7 @@ type Store interface { GetAgentNetworkSettingsBySubdomain(ctx context.Context, lockStrength LockingStrength, subdomain string) (*agentNetworkTypes.Settings, error) SaveAgentNetworkSettings(ctx context.Context, settings *agentNetworkTypes.Settings) error CreateAgentNetworkSettings(ctx context.Context, settings *agentNetworkTypes.Settings) error + SetAgentNetworkServingProxyAddress(ctx context.Context, accountID, address string) error IncrementAgentNetworkConsumption(ctx context.Context, accountID string, kind agentNetworkTypes.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time, tokensIn, tokensOut int64, costUSD float64) error IncrementAgentNetworkConsumptionBatch(ctx context.Context, accountID string, keys []agentNetworkTypes.ConsumptionKey, tokensIn, tokensOut int64, costUSD float64) error GetAgentNetworkConsumption(ctx context.Context, lockStrength LockingStrength, accountID string, kind agentNetworkTypes.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time) (*agentNetworkTypes.Consumption, error) diff --git a/management/server/store/store_mock.go b/management/server/store/store_mock.go index b96cd2f08..ed05eca7d 100644 --- a/management/server/store/store_mock.go +++ b/management/server/store/store_mock.go @@ -3666,6 +3666,20 @@ func (mr *MockStoreMockRecorder) SaveUsers(ctx, users interface{}) *gomock.Call return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveUsers", reflect.TypeOf((*MockStore)(nil).SaveUsers), ctx, users) } +// SetAgentNetworkServingProxyAddress mocks base method. +func (m *MockStore) SetAgentNetworkServingProxyAddress(ctx context.Context, accountID, address string) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "SetAgentNetworkServingProxyAddress", ctx, accountID, address) + ret0, _ := ret[0].(error) + return ret0 +} + +// SetAgentNetworkServingProxyAddress indicates an expected call of SetAgentNetworkServingProxyAddress. +func (mr *MockStoreMockRecorder) SetAgentNetworkServingProxyAddress(ctx, accountID, address interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetAgentNetworkServingProxyAddress", reflect.TypeOf((*MockStore)(nil).SetAgentNetworkServingProxyAddress), ctx, accountID, address) +} + // SetFieldEncrypt mocks base method. func (m *MockStore) SetFieldEncrypt(enc *crypt.FieldEncrypt) { m.ctrl.T.Helper() diff --git a/management/server/types/account.go b/management/server/types/account.go index 059fc20cf..dbdfd3ed9 100644 --- a/management/server/types/account.go +++ b/management/server/types/account.go @@ -364,8 +364,9 @@ func (a *Account) privateServiceDomainZone(svc *service.Service) string { // explicitly: it is server config, so it matches neither the serving // proxy's address nor any per-account custom-domain row. Checked first so // the apex stays the zone even once ProxyCluster becomes the tenant - // hostname itself (a private managed proxy), which would otherwise make the - // apex the full hostname and churn the client's zone set on cutover. + // hostname itself (which happens when a dedicated per-account proxy serves + // it), which would otherwise make the apex the full hostname and churn the + // client's zone set when a tenant moves between proxies. if svc.DNSZone != "" && domainFromSuffix(svc.Domain, svc.DNSZone) { return svc.DNSZone }