From 7f3a29d6888c60000de110f4ec4e2cfce9bd6ec5 Mon Sep 17 00:00:00 2001 From: Brad Ison Date: Mon, 3 Aug 2026 20:42:35 +0200 Subject: [PATCH] feat(agentnetwork): record which proxy serves an account Adds ServingProxyAddress to the agent-network settings row and makes the synthesized service advertise it as the proxy that serves the account. Empty -- the value for every existing row -- means the shared proxy named by Cluster serves the account, which is the previous behaviour exactly. This is the only mutable column on the 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. Because mesh DNS peer selection and the mapping-delivery filter both join on the advertised address, that single write is what redirects an account's traffic, with no change to either join. Only the advertised address follows the new column. The placeholder target's ID stays pinned to Cluster: it names a target the router rewrites per request and is not a routing decision. Nothing in this repository writes the column; it is set by whatever process assigns dedicated proxies. A store setter is included so that assignment can be a single scoped update rather than a read-modify-write over the whole row, which would race with unrelated settings changes. --- .../modules/agentnetwork/synthesizer.go | 8 +- .../modules/agentnetwork/synthesizer_test.go | 96 +++++++++++++++++++ .../modules/agentnetwork/types/settings.go | 33 ++++++- .../agentnetwork/types/settings_test.go | 14 +++ .../server/store/sql_store_agentnetwork.go | 18 ++++ management/server/store/store.go | 1 + management/server/store/store_mock.go | 14 +++ management/server/types/account.go | 5 +- 8 files changed, 183 insertions(+), 6 deletions(-) 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 }