Merge origin/main into embedded-vnc

shared/management/client/grpc.go

@@ -33,10 +33,15 @@ const ConnectTimeout = 10 * time.Second
 const healthCheckTimeout = 5 * time.Second
 
 const (
-	// EnvMaxRecvMsgSize overrides the default gRPC max receive message size (4 MB)
+	// EnvMaxRecvMsgSize overrides the default gRPC max receive message size
 	// for the management client connection. Value is in bytes.
 	EnvMaxRecvMsgSize = "NB_MANAGEMENT_GRPC_MAX_MSG_SIZE"
 
+	// defaultMaxRecvMsgSize is the max gRPC receive message size used for the
+	// management client connection when EnvMaxRecvMsgSize is unset or invalid.
+	// It overrides the gRPC library default of 4 MB.
+	defaultMaxRecvMsgSize = 1024 * 1024 * 16
+
 	errMsgMgmtPublicKey    = "failed getting Management Service public key: %s"
 	errMsgNoMgmtConnection = "no connection to management"
 )
@@ -55,6 +60,14 @@ type GrpcClient struct {
 	connStateCallback     ConnStateNotifier
 	connStateCallbackLock sync.RWMutex
 	serverURL             string
+
+	// syncStreamErr holds the last Sync stream error, or nil while the stream
+	// is established and healthy. GetServerKey succeeds even when the peer
+	// cannot sync (e.g. the server returns "settings not found"), so the
+	// health probe must consult this to avoid reporting a healthy management
+	// connection while the Sync stream keeps failing.
+	syncStreamMu  sync.RWMutex
+	syncStreamErr error
 }
 
 type ExposeRequest struct {
@@ -76,22 +89,22 @@ type ExposeResponse struct {
 }
 
 // MaxRecvMsgSize returns the configured max gRPC receive message size from
-// the environment, or 0 if unset (which uses the gRPC default of 4 MB).
+// the environment, or defaultMaxRecvMsgSize (16 MB) if unset or invalid.
 func MaxRecvMsgSize() int {
 	val := os.Getenv(EnvMaxRecvMsgSize)
 	if val == "" {
-		return 0
+		return defaultMaxRecvMsgSize
 	}
 
 	size, err := strconv.Atoi(val)
 	if err != nil {
 		log.Warnf("invalid %s value %q, using default: %v", EnvMaxRecvMsgSize, val, err)
-		return 0
+		return defaultMaxRecvMsgSize
 	}
 
 	if size <= 0 {
 		log.Warnf("invalid %s value %d, must be positive, using default", EnvMaxRecvMsgSize, size)
-		return 0
+		return defaultMaxRecvMsgSize
 	}
 
 	return size
@@ -364,6 +377,8 @@ func (c *GrpcClient) handleSyncStream(ctx context.Context, serverPubKey wgtypes.
 	stream, err := c.connectToSyncStream(ctx, serverPubKey, sysInfo)
 	if err != nil {
 		log.Debugf("failed to open Management Service stream: %s", err)
+		c.notifyDisconnected(err)
+		c.setSyncStreamDisconnected(err)
 		if s, ok := gstatus.FromError(err); ok && s.Code() == codes.PermissionDenied {
 			return backoff.Permanent(err) // unrecoverable error, propagate to the upper layer
 		}
@@ -372,11 +387,13 @@ func (c *GrpcClient) handleSyncStream(ctx context.Context, serverPubKey wgtypes.
 
 	log.Infof("connected to the Management Service stream")
 	c.notifyConnected()
+	c.setSyncStreamConnected()
 
 	// blocking until error
 	err = c.receiveUpdatesEvents(stream, serverPubKey, msgHandler)
 	if err != nil {
 		c.notifyDisconnected(err)
+		c.setSyncStreamDisconnected(err)
 		if ctx.Err() != nil {
 			log.Debugf("management connection context has been canceled, this usually indicates shutdown")
 			return nil
@@ -524,12 +541,19 @@ func (c *GrpcClient) IsHealthy() bool {
 	ctx, cancel := context.WithTimeout(c.ctx, healthCheckTimeout)
 	defer cancel()
 
-	_, err := c.realClient.GetServerKey(ctx, &proto.Empty{})
+	_, err := c.realClient.IsHealthy(ctx, &proto.Empty{})
 	if err != nil {
 		c.notifyDisconnected(err)
 		log.Warnf("health check returned: %s", err)
 		return false
 	}
+
+	if syncErr := c.syncStreamError(); syncErr != nil {
+		c.notifyDisconnected(syncErr)
+		log.Warnf("management transport is up but the Sync stream is unhealthy: %s", syncErr)
+		return false
+	}
+
 	c.notifyConnected()
 	return true
 }
@@ -630,26 +654,14 @@ func (c *GrpcClient) ExtendAuthSession(sysInfo *system.Info, jwtToken string) (*
 		return nil, err
 	}
 
-	var resp *proto.EncryptedMessage
-	operation := func() error {
-		mgmCtx, cancel := context.WithTimeout(context.Background(), ConnectTimeout)
-		defer cancel()
+	mgmCtx, cancel := context.WithTimeout(c.ctx, ConnectTimeout)
+	defer cancel()
 
-		var err error
-		resp, err = c.realClient.ExtendAuthSession(mgmCtx, &proto.EncryptedMessage{
-			WgPubKey: c.key.PublicKey().String(),
-			Body:     reqBody,
-		})
-		if err != nil {
-			if s, ok := gstatus.FromError(err); ok && s.Code() == codes.Canceled {
-				return err
-			}
-			return backoff.Permanent(err)
-		}
-		return nil
-	}
-
-	if err := backoff.Retry(operation, nbgrpc.Backoff(c.ctx)); err != nil {
+	resp, err := c.realClient.ExtendAuthSession(mgmCtx, &proto.EncryptedMessage{
+		WgPubKey: c.key.PublicKey().String(),
+		Body:     reqBody,
+	})
+	if err != nil {
 		log.Errorf("failed to extend auth session on Management Service: %v", err)
 		return nil, err
 	}
@@ -771,6 +783,24 @@ func (c *GrpcClient) SyncMeta(sysInfo *system.Info) error {
 	return err
 }
 
+func (c *GrpcClient) setSyncStreamConnected() {
+	c.syncStreamMu.Lock()
+	defer c.syncStreamMu.Unlock()
+	c.syncStreamErr = nil
+}
+
+func (c *GrpcClient) setSyncStreamDisconnected(err error) {
+	c.syncStreamMu.Lock()
+	defer c.syncStreamMu.Unlock()
+	c.syncStreamErr = err
+}
+
+func (c *GrpcClient) syncStreamError() error {
+	c.syncStreamMu.RLock()
+	defer c.syncStreamMu.RUnlock()
+	return c.syncStreamErr
+}
+
 func (c *GrpcClient) notifyDisconnected(err error) {
 	c.connStateCallbackLock.RLock()
 	defer c.connStateCallbackLock.RUnlock()
@@ -995,8 +1025,6 @@ func infoToMetaData(info *system.Info) *proto.PeerSystemMeta {
 			BlockInbound:        info.BlockInbound,
 			DisableIPv6:         info.DisableIPv6,
 
-			LazyConnectionEnabled: info.LazyConnectionEnabled,
-
 			DisableSSHAuth: info.DisableSSHAuth,
 		},
 

shared/management/client/grpc_test.go

@@ -21,11 +21,11 @@ func TestMaxRecvMsgSize(t *testing.T) {
 		envValue string
 		expected int
 	}{
-		{name: "unset returns 0", envValue: "", expected: 0},
+		{name: "unset returns default", envValue: "", expected: defaultMaxRecvMsgSize},
 		{name: "valid value", envValue: "10485760", expected: 10485760},
-		{name: "non-numeric returns 0", envValue: "abc", expected: 0},
-		{name: "negative returns 0", envValue: "-1", expected: 0},
-		{name: "zero returns 0", envValue: "0", expected: 0},
+		{name: "non-numeric returns default", envValue: "abc", expected: defaultMaxRecvMsgSize},
+		{name: "negative returns default", envValue: "-1", expected: defaultMaxRecvMsgSize},
+		{name: "zero returns default", envValue: "0", expected: defaultMaxRecvMsgSize},
 	}
 
 	for _, tt := range tests {

shared/management/proto/management.proto

@@ -314,6 +314,8 @@ message NetbirdConfig {
   RelayConfig relay = 4;
 
   FlowConfig flow = 5;
+
+  MetricsConfig metrics = 6;
 }
 
 // HostConfig describes connection properties of some server (e.g. STUN, Signal, Management)
@@ -352,6 +354,10 @@ message FlowConfig {
   bool dnsCollection = 8;
 }
 
+message MetricsConfig {
+  bool enabled = 1;
+}
+
 // JWTConfig represents JWT authentication configuration for validating tokens.
 message JWTConfig {
   string issuer = 1;

shared/management/proto/proxy_service.proto

@@ -43,6 +43,18 @@ service ProxyService {
   // issue a session cookie without redirecting through the OIDC flow.
   // Mirrors ValidateSession's response shape.
   rpc ValidateTunnelPeer(ValidateTunnelPeerRequest) returns (ValidateTunnelPeerResponse);
+
+  // CheckLLMPolicyLimits is the pre-flight RPC the proxy calls before each
+  // LLM request. Management runs the per-policy headroom selection across
+  // every policy authorising the caller's user / groups for the resolved
+  // provider and returns the chosen attribution policy + group, or a deny
+  // when no applicable policy has headroom > 0.
+  rpc CheckLLMPolicyLimits(CheckLLMPolicyLimitsRequest) returns (CheckLLMPolicyLimitsResponse);
+
+  // RecordLLMUsage is the post-flight RPC the proxy calls after the upstream
+  // returns. Increments the per-(dimension, window) counters for the
+  // attribution policy chosen by CheckLLMPolicyLimits.
+  rpc RecordLLMUsage(RecordLLMUsageRequest) returns (RecordLLMUsageResponse);
 }
 
 // ProxyCapabilities describes what a proxy can handle.
@@ -107,6 +119,59 @@ message PathTargetOptions {
   // reachable without WireGuard (public APIs, LAN services, localhost
   // sidecars). Defaults to false — embedded client is the standard path.
   bool direct_upstream = 7;
+  // Proxy clamps to [0, proxy-wide max (1 MiB)] at apply time. Agent-network
+  // synthesized targets only; private services leave these zero.
+  int64 capture_max_request_bytes = 8;
+  // Proxy clamps to [0, proxy-wide max (1 MiB)] at apply time.
+  int64 capture_max_response_bytes = 9;
+  // Content types eligible for body capture (e.g. "application/json").
+  repeated string capture_content_types = 10;
+  // Per-target middleware configurations populated by the agent-network
+  // synthesizer. Validated and clamped by the proxy at apply time.
+  repeated MiddlewareConfig middlewares = 11;
+  // When true, the proxy stamps agent_network=true on access-log entries
+  // for this target so management routes them to the agent-network log
+  // surface.
+  bool agent_network = 12;
+  // When true, the proxy suppresses the per-request access-log emission for
+  // this target. Defaults false to preserve existing access-log behavior for
+  // every non-agent-network target. The agent-network synth target sets this
+  // true only when the account's EnableLogCollection toggle is off.
+  bool disable_access_log = 13;
+}
+
+// MiddlewareSlot identifies where in the request lifecycle a middleware
+// runs. Mirrors proxy/internal/middleware.Slot.
+enum MiddlewareSlot {
+  MIDDLEWARE_SLOT_UNSPECIFIED = 0;
+  MIDDLEWARE_SLOT_ON_REQUEST = 1;
+  MIDDLEWARE_SLOT_ON_RESPONSE = 2;
+  MIDDLEWARE_SLOT_TERMINAL = 3;
+}
+
+// MiddlewareConfig is the per-target configuration for a single middleware.
+// The proxy validates every incoming MiddlewareConfig at apply time:
+// unknown ids are rejected, timeout is clamped to [10ms, 5s], and the
+// declared slot must match the registered middleware's slot.
+message MiddlewareConfig {
+  // Middleware id; must match the proxy-local compiled-in registry.
+  string id = 1;
+  bool enabled = 2;
+  MiddlewareSlot slot = 3;
+  // Free-form JSON unmarshalled by the middleware factory into its own typed
+  // config struct. Empty / null / {} are valid (zero-value config).
+  bytes config_json = 4;
+  enum FailMode {
+    FAIL_OPEN = 0;
+    FAIL_CLOSED = 1;
+  }
+  FailMode fail_mode = 5;
+  // Clamped to [10ms, 5s] at apply time; zero → 500ms default.
+  google.protobuf.Duration timeout = 6;
+  // When true, the middleware may mutate request headers or body (subject to
+  // policy). Honoured only when the implementation also declares
+  // MutationsSupported.
+  bool can_mutate = 7;
 }
 
 message PathMapping {
@@ -190,6 +255,10 @@ message AccessLog {
   string protocol = 16;
   // Extra key-value metadata for the access log entry (e.g. crowdsec_verdict, scenario).
   map<string, string> metadata = 17;
+  // When true, the entry was emitted by an agent-network synth service.
+  // Management routes these to the agent-network access-log surface instead
+  // of the standard service log.
+  bool agent_network = 18;
 }
 
 message AuthenticateRequest {
@@ -376,3 +445,59 @@ message SyncMappingsResponse {
   bool initial_sync_complete = 2;
 }
 
+// CheckLLMPolicyLimitsRequest carries the resolved caller identity and the
+// upstream provider already chosen by llm_router. Management computes which
+// policies authorise the request, picks the one with the most remaining
+// headroom, and returns the attribution decision.
+message CheckLLMPolicyLimitsRequest {
+  // account_id is the netbird account the request belongs to.
+  string account_id = 1;
+  // user_id is the netbird user id of the caller. May be empty when the
+  // principal is a tunnel-peer that isn't bound to a user; group membership
+  // still gates the request in that case.
+  string user_id = 2;
+  // group_ids is the caller's full group membership at request time.
+  repeated string group_ids = 3;
+  // provider_id is the agent-network provider record id chosen by llm_router.
+  string provider_id = 4;
+  // model is the upstream model identifier extracted from the request body.
+  string model = 5;
+}
+
+// CheckLLMPolicyLimitsResponse is management's allow-or-deny decision for a
+// pre-flight check.
+message CheckLLMPolicyLimitsResponse {
+  // decision is "allow" or "deny".
+  string decision = 1;
+  // selected_policy_id names the policy that paid for this request.
+  string selected_policy_id = 2;
+  // attribution_group_id is the source group the request booked against.
+  string attribution_group_id = 3;
+  // window_seconds is the cap window length the selected policy uses.
+  int64 window_seconds = 4;
+  // deny_code is set on decision="deny" with a stable label.
+  string deny_code = 5;
+  // deny_reason is a short human-readable explanation paired with deny_code.
+  string deny_reason = 6;
+}
+
+// RecordLLMUsageRequest is the post-flight increment the proxy posts after
+// the upstream call. Counters are keyed on (account, dimension, window).
+message RecordLLMUsageRequest {
+  string account_id = 1;
+  string user_id = 2;
+  // group_id is the selected policy's attribution group, recorded against the
+  // policy window (window_seconds).
+  string group_id = 3;
+  int64 window_seconds = 4;
+  int64 tokens_input = 5;
+  int64 tokens_output = 6;
+  double cost_usd = 7;
+  // group_ids is the caller's full group membership, used to fan the same
+  // usage out to every applicable account-level budget rule's own window.
+  repeated string group_ids = 8;
+}
+
+message RecordLLMUsageResponse {
+}
+

shared/management/proto/proxy_service_grpc.pb.go

@@ -43,6 +43,16 @@ type ProxyServiceClient interface {
 	// issue a session cookie without redirecting through the OIDC flow.
 	// Mirrors ValidateSession's response shape.
 	ValidateTunnelPeer(ctx context.Context, in *ValidateTunnelPeerRequest, opts ...grpc.CallOption) (*ValidateTunnelPeerResponse, error)
+	// CheckLLMPolicyLimits is the pre-flight RPC the proxy calls before each
+	// LLM request. Management runs the per-policy headroom selection across
+	// every policy authorising the caller's user / groups for the resolved
+	// provider and returns the chosen attribution policy + group, or a deny
+	// when no applicable policy has headroom > 0.
+	CheckLLMPolicyLimits(ctx context.Context, in *CheckLLMPolicyLimitsRequest, opts ...grpc.CallOption) (*CheckLLMPolicyLimitsResponse, error)
+	// RecordLLMUsage is the post-flight RPC the proxy calls after the upstream
+	// returns. Increments the per-(dimension, window) counters for the
+	// attribution policy chosen by CheckLLMPolicyLimits.
+	RecordLLMUsage(ctx context.Context, in *RecordLLMUsageRequest, opts ...grpc.CallOption) (*RecordLLMUsageResponse, error)
 }
 
 type proxyServiceClient struct {
@@ -179,6 +189,24 @@ func (c *proxyServiceClient) ValidateTunnelPeer(ctx context.Context, in *Validat
 	return out, nil
 }
 
+func (c *proxyServiceClient) CheckLLMPolicyLimits(ctx context.Context, in *CheckLLMPolicyLimitsRequest, opts ...grpc.CallOption) (*CheckLLMPolicyLimitsResponse, error) {
+	out := new(CheckLLMPolicyLimitsResponse)
+	err := c.cc.Invoke(ctx, "/management.ProxyService/CheckLLMPolicyLimits", in, out, opts...)
+	if err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
+func (c *proxyServiceClient) RecordLLMUsage(ctx context.Context, in *RecordLLMUsageRequest, opts ...grpc.CallOption) (*RecordLLMUsageResponse, error) {
+	out := new(RecordLLMUsageResponse)
+	err := c.cc.Invoke(ctx, "/management.ProxyService/RecordLLMUsage", in, out, opts...)
+	if err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
 // ProxyServiceServer is the server API for ProxyService service.
 // All implementations must embed UnimplementedProxyServiceServer
 // for forward compatibility
@@ -208,6 +236,16 @@ type ProxyServiceServer interface {
 	// issue a session cookie without redirecting through the OIDC flow.
 	// Mirrors ValidateSession's response shape.
 	ValidateTunnelPeer(context.Context, *ValidateTunnelPeerRequest) (*ValidateTunnelPeerResponse, error)
+	// CheckLLMPolicyLimits is the pre-flight RPC the proxy calls before each
+	// LLM request. Management runs the per-policy headroom selection across
+	// every policy authorising the caller's user / groups for the resolved
+	// provider and returns the chosen attribution policy + group, or a deny
+	// when no applicable policy has headroom > 0.
+	CheckLLMPolicyLimits(context.Context, *CheckLLMPolicyLimitsRequest) (*CheckLLMPolicyLimitsResponse, error)
+	// RecordLLMUsage is the post-flight RPC the proxy calls after the upstream
+	// returns. Increments the per-(dimension, window) counters for the
+	// attribution policy chosen by CheckLLMPolicyLimits.
+	RecordLLMUsage(context.Context, *RecordLLMUsageRequest) (*RecordLLMUsageResponse, error)
 	mustEmbedUnimplementedProxyServiceServer()
 }
 
@@ -242,6 +280,12 @@ func (UnimplementedProxyServiceServer) ValidateSession(context.Context, *Validat
 func (UnimplementedProxyServiceServer) ValidateTunnelPeer(context.Context, *ValidateTunnelPeerRequest) (*ValidateTunnelPeerResponse, error) {
 	return nil, status.Errorf(codes.Unimplemented, "method ValidateTunnelPeer not implemented")
 }
+func (UnimplementedProxyServiceServer) CheckLLMPolicyLimits(context.Context, *CheckLLMPolicyLimitsRequest) (*CheckLLMPolicyLimitsResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method CheckLLMPolicyLimits not implemented")
+}
+func (UnimplementedProxyServiceServer) RecordLLMUsage(context.Context, *RecordLLMUsageRequest) (*RecordLLMUsageResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method RecordLLMUsage not implemented")
+}
 func (UnimplementedProxyServiceServer) mustEmbedUnimplementedProxyServiceServer() {}
 
 // UnsafeProxyServiceServer may be embedded to opt out of forward compatibility for this service.
@@ -428,6 +472,42 @@ func _ProxyService_ValidateTunnelPeer_Handler(srv interface{}, ctx context.Conte
 	return interceptor(ctx, in, info, handler)
 }
 
+func _ProxyService_CheckLLMPolicyLimits_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(CheckLLMPolicyLimitsRequest)
+	if err := dec(in); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(ProxyServiceServer).CheckLLMPolicyLimits(ctx, in)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: "/management.ProxyService/CheckLLMPolicyLimits",
+	}
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return srv.(ProxyServiceServer).CheckLLMPolicyLimits(ctx, req.(*CheckLLMPolicyLimitsRequest))
+	}
+	return interceptor(ctx, in, info, handler)
+}
+
+func _ProxyService_RecordLLMUsage_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(RecordLLMUsageRequest)
+	if err := dec(in); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(ProxyServiceServer).RecordLLMUsage(ctx, in)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: "/management.ProxyService/RecordLLMUsage",
+	}
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return srv.(ProxyServiceServer).RecordLLMUsage(ctx, req.(*RecordLLMUsageRequest))
+	}
+	return interceptor(ctx, in, info, handler)
+}
+
 // ProxyService_ServiceDesc is the grpc.ServiceDesc for ProxyService service.
 // It's only intended for direct use with grpc.RegisterService,
 // and not to be introspected or modified (even as a copy)
@@ -463,6 +543,14 @@ var ProxyService_ServiceDesc = grpc.ServiceDesc{
 			MethodName: "ValidateTunnelPeer",
 			Handler:    _ProxyService_ValidateTunnelPeer_Handler,
 		},
+		{
+			MethodName: "CheckLLMPolicyLimits",
+			Handler:    _ProxyService_CheckLLMPolicyLimits_Handler,
+		},
+		{
+			MethodName: "RecordLLMUsage",
+			Handler:    _ProxyService_RecordLLMUsage_Handler,
+		},
 	},
 	Streams: []grpc.StreamDesc{
 		{

shared/management/status/error.go

@@ -219,6 +219,26 @@ func NewNetworkResourceNotFoundError(resourceID string) error {
 	return Errorf(NotFound, "network resource: %s not found", resourceID)
 }
 
+// NewAgentNetworkProviderNotFoundError creates a new Error with NotFound type for a missing Agent Network provider.
+func NewAgentNetworkProviderNotFoundError(providerID string) error {
+	return Errorf(NotFound, "agent network provider: %s not found", providerID)
+}
+
+// NewAgentNetworkPolicyNotFoundError creates a new Error with NotFound type for a missing Agent Network policy.
+func NewAgentNetworkPolicyNotFoundError(policyID string) error {
+	return Errorf(NotFound, "agent network policy: %s not found", policyID)
+}
+
+// NewAgentNetworkGuardrailNotFoundError creates a new Error with NotFound type for a missing Agent Network guardrail.
+func NewAgentNetworkGuardrailNotFoundError(guardrailID string) error {
+	return Errorf(NotFound, "agent network guardrail: %s not found", guardrailID)
+}
+
+// NewAgentNetworkBudgetRuleNotFoundError creates a new Error with NotFound type for a missing Agent Network budget rule.
+func NewAgentNetworkBudgetRuleNotFoundError(ruleID string) error {
+	return Errorf(NotFound, "agent network budget rule: %s not found", ruleID)
+}
+
 // NewPermissionDeniedError creates a new Error with PermissionDenied type for a permission denied error.
 func NewPermissionDeniedError() error {
 	return Errorf(PermissionDenied, "permission denied")