diff --git a/chatserver/server.py b/chatserver/server.py
index cda46f7..68f5483 100644
--- a/chatserver/server.py
+++ b/chatserver/server.py
@@ -144,24 +144,39 @@ def handle_command(username, chat_room, command, sock=None):
             target_username = parts[1]
             dm_message = parts[2]
             timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+
+            # Check if the target is a socket-based client
             target_sock = None
             for s, info in client_info.items():
-                if info.get("username") == target_username:
+                if not info.get("is_websocket") and info.get("username") == target_username:
                     target_sock = s
                     break
+
+            # Check if the target is a WebSocket client
+            target_websocket_sid = None
+            for sid, info in client_info.items():
+                if info.get("is_websocket") and info.get("username") == target_username:
+                    target_websocket_sid = sid
+                    break
+
             if target_sock:
+                # Send the direct message to the socket-based client
                 sender_username = username
                 target_sock.send(f"[{timestamp}] DM from {sender_username}: {dm_message}\n".encode('utf-8'))
                 if sock:
                     sock.send(f"[{timestamp}] Direct message sent.\n".encode('utf-8'))
-                else:
-                    socketio.emit('message', {
-                        'username': sender_username,
-                        'message': dm_message,
-                        'timestamp': timestamp,
-                        'chat_room': 'Direct Message'
-                    }, room=chat_room)
+            elif target_websocket_sid:
+                # Send the direct message to the WebSocket client
+                socketio.emit('message', {
+                    'username': username,
+                    'message': dm_message,
+                    'timestamp': timestamp,
+                    'chat_room': 'Direct Message'
+                }, to=target_websocket_sid)
+                if sock:
+                    sock.send(f"[{timestamp}] Direct message sent.\n".encode('utf-8'))
             else:
+                # Target user not found
                 error_message = "User not found."
                 if sock:
                     sock.send(f"\n{error_message}\n".encode('utf-8'))
@@ -237,6 +252,10 @@ def get_history():
 @socketio.on('connect')
 def handle_connect():
     print("WebSocket client connected")
+    # Assign a unique ID to the WebSocket client
+    user_id = str(uuid.uuid4())
+    client_info[request.sid] = {"user_id": user_id, "username": None, "is_websocket": True}
+    print(f"WebSocket client registered with ID: {user_id}")
 
 @socketio.on('disconnect')
 def handle_disconnect():
@@ -252,6 +271,10 @@ def handle_join(data):
         emit('error', {'message': 'Username is required'})
         return
 
+    # Update client info
+    client_info[request.sid]["username"] = username
+    client_info[request.sid]["chat_room"] = chat_room
+
     join_room(chat_room)
     emit('message', {'username': 'Server', 'message': f"{username} has joined the chat room '{chat_room}'."}, room=chat_room)