[client] Add sleep state tracking to handle wakeup/sleep events reliably (#4894)

Adds a new NotifyOSLifecycle RPC and server handler to centralize OS sleep/wake handling, introduces Server.sleepTriggeredDown for coordination, updates client UI to call the new RPC, and adjusts the internal sleep event enum zero-value semantics.

client/internal/sleep/service.go

@@ -1,8 +1,9 @@
 package sleep
 
 var (
-	EventTypeSleep  EventType = 0
-	EventTypeWakeUp EventType = 1
+	EventTypeUnknown EventType = 0
+	EventTypeSleep   EventType = 1
+	EventTypeWakeUp  EventType = 2
 )
 
 type EventType int

client/proto/daemon.proto

@@ -24,7 +24,7 @@ service DaemonService {
   // Status of the service.
   rpc Status(StatusRequest) returns (StatusResponse) {}
 
-  // Down engine work in the daemon.
+  // Down stops engine work in the daemon.
   rpc Down(DownRequest) returns (DownResponse) {}
 
   // GetConfig of the daemon.
@@ -93,9 +93,26 @@ service DaemonService {
 
   // WaitJWTToken waits for JWT authentication completion
   rpc WaitJWTToken(WaitJWTTokenRequest) returns (WaitJWTTokenResponse) {}
+
+  rpc NotifyOSLifecycle(OSLifecycleRequest) returns(OSLifecycleResponse) {}
 }
 
 
+
+message OSLifecycleRequest {
+  // avoid collision with loglevel enum
+  enum CycleType {
+    UNKNOWN = 0;
+    SLEEP = 1;
+    WAKEUP = 2;
+  }
+
+  CycleType type = 1;
+}
+
+message OSLifecycleResponse {}
+
+
 message LoginRequest {
   // setupKey netbird setup key.
   string setupKey = 1;

client/proto/daemon_grpc.pb.go

@@ -27,7 +27,7 @@ type DaemonServiceClient interface {
 	Up(ctx context.Context, in *UpRequest, opts ...grpc.CallOption) (*UpResponse, error)
 	// Status of the service.
 	Status(ctx context.Context, in *StatusRequest, opts ...grpc.CallOption) (*StatusResponse, error)
-	// Down engine work in the daemon.
+	// Down stops engine work in the daemon.
 	Down(ctx context.Context, in *DownRequest, opts ...grpc.CallOption) (*DownResponse, error)
 	// GetConfig of the daemon.
 	GetConfig(ctx context.Context, in *GetConfigRequest, opts ...grpc.CallOption) (*GetConfigResponse, error)
@@ -70,6 +70,7 @@ type DaemonServiceClient interface {
 	RequestJWTAuth(ctx context.Context, in *RequestJWTAuthRequest, opts ...grpc.CallOption) (*RequestJWTAuthResponse, error)
 	// WaitJWTToken waits for JWT authentication completion
 	WaitJWTToken(ctx context.Context, in *WaitJWTTokenRequest, opts ...grpc.CallOption) (*WaitJWTTokenResponse, error)
+	NotifyOSLifecycle(ctx context.Context, in *OSLifecycleRequest, opts ...grpc.CallOption) (*OSLifecycleResponse, error)
 }
 
 type daemonServiceClient struct {
@@ -382,6 +383,15 @@ func (c *daemonServiceClient) WaitJWTToken(ctx context.Context, in *WaitJWTToken
 	return out, nil
 }
 
+func (c *daemonServiceClient) NotifyOSLifecycle(ctx context.Context, in *OSLifecycleRequest, opts ...grpc.CallOption) (*OSLifecycleResponse, error) {
+	out := new(OSLifecycleResponse)
+	err := c.cc.Invoke(ctx, "/daemon.DaemonService/NotifyOSLifecycle", in, out, opts...)
+	if err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
 // DaemonServiceServer is the server API for DaemonService service.
 // All implementations must embed UnimplementedDaemonServiceServer
 // for forward compatibility
@@ -395,7 +405,7 @@ type DaemonServiceServer interface {
 	Up(context.Context, *UpRequest) (*UpResponse, error)
 	// Status of the service.
 	Status(context.Context, *StatusRequest) (*StatusResponse, error)
-	// Down engine work in the daemon.
+	// Down stops engine work in the daemon.
 	Down(context.Context, *DownRequest) (*DownResponse, error)
 	// GetConfig of the daemon.
 	GetConfig(context.Context, *GetConfigRequest) (*GetConfigResponse, error)
@@ -438,6 +448,7 @@ type DaemonServiceServer interface {
 	RequestJWTAuth(context.Context, *RequestJWTAuthRequest) (*RequestJWTAuthResponse, error)
 	// WaitJWTToken waits for JWT authentication completion
 	WaitJWTToken(context.Context, *WaitJWTTokenRequest) (*WaitJWTTokenResponse, error)
+	NotifyOSLifecycle(context.Context, *OSLifecycleRequest) (*OSLifecycleResponse, error)
 	mustEmbedUnimplementedDaemonServiceServer()
 }
 
@@ -538,6 +549,9 @@ func (UnimplementedDaemonServiceServer) RequestJWTAuth(context.Context, *Request
 func (UnimplementedDaemonServiceServer) WaitJWTToken(context.Context, *WaitJWTTokenRequest) (*WaitJWTTokenResponse, error) {
 	return nil, status.Errorf(codes.Unimplemented, "method WaitJWTToken not implemented")
 }
+func (UnimplementedDaemonServiceServer) NotifyOSLifecycle(context.Context, *OSLifecycleRequest) (*OSLifecycleResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method NotifyOSLifecycle not implemented")
+}
 func (UnimplementedDaemonServiceServer) mustEmbedUnimplementedDaemonServiceServer() {}
 
 // UnsafeDaemonServiceServer may be embedded to opt out of forward compatibility for this service.
@@ -1112,6 +1126,24 @@ func _DaemonService_WaitJWTToken_Handler(srv interface{}, ctx context.Context, d
 	return interceptor(ctx, in, info, handler)
 }
 
+func _DaemonService_NotifyOSLifecycle_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(OSLifecycleRequest)
+	if err := dec(in); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(DaemonServiceServer).NotifyOSLifecycle(ctx, in)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: "/daemon.DaemonService/NotifyOSLifecycle",
+	}
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return srv.(DaemonServiceServer).NotifyOSLifecycle(ctx, req.(*OSLifecycleRequest))
+	}
+	return interceptor(ctx, in, info, handler)
+}
+
 // DaemonService_ServiceDesc is the grpc.ServiceDesc for DaemonService service.
 // It's only intended for direct use with grpc.RegisterService,
 // and not to be introspected or modified (even as a copy)
@@ -1239,6 +1271,10 @@ var DaemonService_ServiceDesc = grpc.ServiceDesc{
 			MethodName: "WaitJWTToken",
 			Handler:    _DaemonService_WaitJWTToken_Handler,
 		},
+		{
+			MethodName: "NotifyOSLifecycle",
+			Handler:    _DaemonService_NotifyOSLifecycle_Handler,
+		},
 	},
 	Streams: []grpc.StreamDesc{
 		{

client/server/lifecycle.go

@@ -0,0 +1,77 @@
+package server
+
+import (
+	"context"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/client/internal"
+	"github.com/netbirdio/netbird/client/proto"
+)
+
+// NotifyOSLifecycle handles operating system lifecycle events by executing appropriate logic based on the request type.
+func (s *Server) NotifyOSLifecycle(callerCtx context.Context, req *proto.OSLifecycleRequest) (*proto.OSLifecycleResponse, error) {
+	switch req.GetType() {
+	case proto.OSLifecycleRequest_WAKEUP:
+		return s.handleWakeUp(callerCtx)
+	case proto.OSLifecycleRequest_SLEEP:
+		return s.handleSleep(callerCtx)
+	default:
+		log.Errorf("unknown OSLifecycleRequest type: %v", req.GetType())
+	}
+	return &proto.OSLifecycleResponse{}, nil
+}
+
+// handleWakeUp processes a wake-up event by triggering the Up command if the system was previously put to sleep.
+// It resets the sleep state and logs the process. Returns a response or an error if the Up command fails.
+func (s *Server) handleWakeUp(callerCtx context.Context) (*proto.OSLifecycleResponse, error) {
+	if !s.sleepTriggeredDown.Load() {
+		log.Info("skipping up because wasn't sleep down")
+		return &proto.OSLifecycleResponse{}, nil
+	}
+
+	// avoid other wakeup runs if sleep didn't make the computer sleep
+	s.sleepTriggeredDown.Store(false)
+
+	log.Info("running up after wake up")
+	_, err := s.Up(callerCtx, &proto.UpRequest{})
+	if err != nil {
+		log.Errorf("running up failed: %v", err)
+		return &proto.OSLifecycleResponse{}, err
+	}
+
+	log.Info("running up command executed successfully")
+	return &proto.OSLifecycleResponse{}, nil
+}
+
+// handleSleep handles the sleep event by initiating a "down" sequence if the system is in a connected or connecting state.
+func (s *Server) handleSleep(callerCtx context.Context) (*proto.OSLifecycleResponse, error) {
+	s.mutex.Lock()
+
+	state := internal.CtxGetState(s.rootCtx)
+	status, err := state.Status()
+	if err != nil {
+		s.mutex.Unlock()
+		return &proto.OSLifecycleResponse{}, err
+	}
+
+	if status != internal.StatusConnecting && status != internal.StatusConnected {
+		log.Infof("skipping setting the agent down because status is %s", status)
+		s.mutex.Unlock()
+		return &proto.OSLifecycleResponse{}, nil
+	}
+	s.mutex.Unlock()
+
+	log.Info("running down after system started sleeping")
+
+	_, err = s.Down(callerCtx, &proto.DownRequest{})
+	if err != nil {
+		log.Errorf("running down failed: %v", err)
+		return &proto.OSLifecycleResponse{}, err
+	}
+
+	s.sleepTriggeredDown.Store(true)
+
+	log.Info("running down executed successfully")
+	return &proto.OSLifecycleResponse{}, nil
+}

client/server/lifecycle_test.go

@@ -0,0 +1,219 @@
+package server
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/client/internal"
+	"github.com/netbirdio/netbird/client/internal/peer"
+	"github.com/netbirdio/netbird/client/proto"
+)
+
+func newTestServer() *Server {
+	ctx := internal.CtxInitState(context.Background())
+	return &Server{
+		rootCtx:        ctx,
+		statusRecorder: peer.NewRecorder(""),
+	}
+}
+
+func TestNotifyOSLifecycle_WakeUp_SkipsWhenNotSleepTriggered(t *testing.T) {
+	s := newTestServer()
+
+	// sleepTriggeredDown is false by default
+	assert.False(t, s.sleepTriggeredDown.Load())
+
+	resp, err := s.NotifyOSLifecycle(context.Background(), &proto.OSLifecycleRequest{
+		Type: proto.OSLifecycleRequest_WAKEUP,
+	})
+
+	require.NoError(t, err)
+	require.NotNil(t, resp)
+	assert.False(t, s.sleepTriggeredDown.Load(), "flag should remain false")
+}
+
+func TestNotifyOSLifecycle_Sleep_SkipsWhenStatusIdle(t *testing.T) {
+	s := newTestServer()
+
+	state := internal.CtxGetState(s.rootCtx)
+	state.Set(internal.StatusIdle)
+
+	resp, err := s.NotifyOSLifecycle(context.Background(), &proto.OSLifecycleRequest{
+		Type: proto.OSLifecycleRequest_SLEEP,
+	})
+
+	require.NoError(t, err)
+	require.NotNil(t, resp)
+	assert.False(t, s.sleepTriggeredDown.Load(), "flag should remain false when status is Idle")
+}
+
+func TestNotifyOSLifecycle_Sleep_SkipsWhenStatusNeedsLogin(t *testing.T) {
+	s := newTestServer()
+
+	state := internal.CtxGetState(s.rootCtx)
+	state.Set(internal.StatusNeedsLogin)
+
+	resp, err := s.NotifyOSLifecycle(context.Background(), &proto.OSLifecycleRequest{
+		Type: proto.OSLifecycleRequest_SLEEP,
+	})
+
+	require.NoError(t, err)
+	require.NotNil(t, resp)
+	assert.False(t, s.sleepTriggeredDown.Load(), "flag should remain false when status is NeedsLogin")
+}
+
+func TestNotifyOSLifecycle_Sleep_SetsFlag_WhenConnecting(t *testing.T) {
+	s := newTestServer()
+
+	state := internal.CtxGetState(s.rootCtx)
+	state.Set(internal.StatusConnecting)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	s.actCancel = cancel
+
+	resp, err := s.NotifyOSLifecycle(ctx, &proto.OSLifecycleRequest{
+		Type: proto.OSLifecycleRequest_SLEEP,
+	})
+
+	require.NoError(t, err)
+	assert.NotNil(t, resp, "handleSleep returns not nil response on success")
+	assert.True(t, s.sleepTriggeredDown.Load(), "flag should be set after sleep when connecting")
+}
+
+func TestNotifyOSLifecycle_Sleep_SetsFlag_WhenConnected(t *testing.T) {
+	s := newTestServer()
+
+	state := internal.CtxGetState(s.rootCtx)
+	state.Set(internal.StatusConnected)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	s.actCancel = cancel
+
+	resp, err := s.NotifyOSLifecycle(ctx, &proto.OSLifecycleRequest{
+		Type: proto.OSLifecycleRequest_SLEEP,
+	})
+
+	require.NoError(t, err)
+	assert.NotNil(t, resp, "handleSleep returns not nil response on success")
+	assert.True(t, s.sleepTriggeredDown.Load(), "flag should be set after sleep when connected")
+}
+
+func TestNotifyOSLifecycle_WakeUp_ResetsFlag(t *testing.T) {
+	s := newTestServer()
+
+	// Manually set the flag to simulate prior sleep down
+	s.sleepTriggeredDown.Store(true)
+
+	// WakeUp will try to call Up which fails without proper setup, but flag should reset first
+	_, _ = s.NotifyOSLifecycle(context.Background(), &proto.OSLifecycleRequest{
+		Type: proto.OSLifecycleRequest_WAKEUP,
+	})
+
+	assert.False(t, s.sleepTriggeredDown.Load(), "flag should be reset after WakeUp attempt")
+}
+
+func TestNotifyOSLifecycle_MultipleWakeUpCalls(t *testing.T) {
+	s := newTestServer()
+
+	// First wakeup without prior sleep - should be no-op
+	resp, err := s.NotifyOSLifecycle(context.Background(), &proto.OSLifecycleRequest{
+		Type: proto.OSLifecycleRequest_WAKEUP,
+	})
+	require.NoError(t, err)
+	require.NotNil(t, resp)
+	assert.False(t, s.sleepTriggeredDown.Load())
+
+	// Simulate prior sleep
+	s.sleepTriggeredDown.Store(true)
+
+	// First wakeup after sleep - should reset flag
+	_, _ = s.NotifyOSLifecycle(context.Background(), &proto.OSLifecycleRequest{
+		Type: proto.OSLifecycleRequest_WAKEUP,
+	})
+	assert.False(t, s.sleepTriggeredDown.Load())
+
+	// Second wakeup - should be no-op
+	resp, err = s.NotifyOSLifecycle(context.Background(), &proto.OSLifecycleRequest{
+		Type: proto.OSLifecycleRequest_WAKEUP,
+	})
+	require.NoError(t, err)
+	require.NotNil(t, resp)
+	assert.False(t, s.sleepTriggeredDown.Load())
+}
+
+func TestHandleWakeUp_SkipsWhenFlagFalse(t *testing.T) {
+	s := newTestServer()
+
+	resp, err := s.handleWakeUp(context.Background())
+
+	require.NoError(t, err)
+	require.NotNil(t, resp)
+}
+
+func TestHandleWakeUp_ResetsFlagBeforeUp(t *testing.T) {
+	s := newTestServer()
+	s.sleepTriggeredDown.Store(true)
+
+	// Even if Up fails, flag should be reset
+	_, _ = s.handleWakeUp(context.Background())
+
+	assert.False(t, s.sleepTriggeredDown.Load(), "flag must be reset before calling Up")
+}
+
+func TestHandleSleep_SkipsForNonActiveStates(t *testing.T) {
+	tests := []struct {
+		name   string
+		status internal.StatusType
+	}{
+		{"Idle", internal.StatusIdle},
+		{"NeedsLogin", internal.StatusNeedsLogin},
+		{"LoginFailed", internal.StatusLoginFailed},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			s := newTestServer()
+			state := internal.CtxGetState(s.rootCtx)
+			state.Set(tt.status)
+
+			resp, err := s.handleSleep(context.Background())
+
+			require.NoError(t, err)
+			require.NotNil(t, resp)
+			assert.False(t, s.sleepTriggeredDown.Load())
+		})
+	}
+}
+
+func TestHandleSleep_ProceedsForActiveStates(t *testing.T) {
+	tests := []struct {
+		name   string
+		status internal.StatusType
+	}{
+		{"Connecting", internal.StatusConnecting},
+		{"Connected", internal.StatusConnected},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			s := newTestServer()
+			state := internal.CtxGetState(s.rootCtx)
+			state.Set(tt.status)
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			s.actCancel = cancel
+
+			resp, err := s.handleSleep(ctx)
+
+			require.NoError(t, err)
+			assert.NotNil(t, resp)
+			assert.True(t, s.sleepTriggeredDown.Load())
+		})
+	}
+}

client/server/server.go

@@ -85,6 +85,9 @@ type Server struct {
 	profilesDisabled       bool
 	updateSettingsDisabled bool
 
+	// sleepTriggeredDown holds a state indicated if the sleep handler triggered the last client down
+	sleepTriggeredDown atomic.Bool
+
 	jwtCache *jwtCache
 }
 

client/ui/client_ui.go

@@ -1196,25 +1196,27 @@ func (s *serviceClient) handleSleepEvents(event sleep.EventType) {
 		return
 	}
 
+	req := &proto.OSLifecycleRequest{}
+
 	switch event {
 	case sleep.EventTypeWakeUp:
 		log.Infof("handle wakeup event: %v", event)
-		_, err = conn.Up(s.ctx, &proto.UpRequest{})
-		if err != nil {
-			log.Errorf("up service: %v", err)
-			return
-		}
-		return
+		req.Type = proto.OSLifecycleRequest_WAKEUP
 	case sleep.EventTypeSleep:
 		log.Infof("handle sleep event: %v", event)
-		_, err = conn.Down(s.ctx, &proto.DownRequest{})
-		if err != nil {
-			log.Errorf("down service: %v", err)
-			return
-		}
+		req.Type = proto.OSLifecycleRequest_SLEEP
+	default:
+		log.Infof("unknown event: %v", event)
+		return
 	}
 
-	log.Info("successfully notified daemon about sleep/wakeup event")
+	_, err = conn.NotifyOSLifecycle(s.ctx, req)
+	if err != nil {
+		log.Errorf("failed to notify daemon about os lifecycle notification: %v", err)
+		return
+	}
+
+	log.Info("successfully notified daemon about os lifecycle")
 }
 
 // setSettingsEnabled enables or disables the settings menu based on the provided state