diff --git a/connections/implementation/mediums/bluetooth_classic.cc b/connections/implementation/mediums/bluetooth_classic.cc
index f38dc278..290defb5 100644
--- a/connections/implementation/mediums/bluetooth_classic.cc
+++ b/connections/implementation/mediums/bluetooth_classic.cc
@@ -64,9 +64,7 @@ BluetoothClassic::BluetoothClassic(
     BluetoothRadio& radio, std::unique_ptr<BluetoothClassicMedium> medium)
     : radio_(radio),
       adapter_(radio_.GetBluetoothAdapter()),
-      medium_(std::move(medium)) {
-  is_multiplex_enabled_ = false;
-}
+      medium_(std::move(medium)) {}
 
 BluetoothClassic::~BluetoothClassic() {
   // Destructor is not taking locks, but methods it is calling are.
@@ -382,10 +380,10 @@ ErrorOr<bool> BluetoothClassic::StartAcceptingConnections(
     MultiplexSocket::ListenForIncomingConnection(
         service_id, Medium::BLUETOOTH,
         [&callback](const std::string& listening_service_id,
-                    MediumSocket* virtual_socket) mutable {
+                    std::shared_ptr<MediumSocket> virtual_socket) mutable {
           if (callback) {
             callback(listening_service_id,
-                     *(down_cast<BluetoothSocket*>(virtual_socket)));
+                     *(down_cast<BluetoothSocket*>(virtual_socket.get())));
           }
         });
   }
@@ -415,20 +413,21 @@ ErrorOr<bool> BluetoothClassic::StartAcceptingConnections(
               MultiplexSocket::CreateIncomingSocket(physical_socket_ptr,
                                                     service_id, 0);
 
-          if (multiplex_socket != nullptr &&
-              multiplex_socket->GetVirtualSocket(service_id)) {
-            multiplex_sockets_.emplace(
-                client_socket.GetRemoteDevice().GetAddress(),
-                multiplex_socket);
-            MultiplexSocket::StopListeningForIncomingConnection(
-                service_id, Medium::BLUETOOTH);
-            LOG(INFO) << "Multiplex virtaul socket created for "
-                      << client_socket.GetRemoteDevice().GetName();
-            if (callback) {
-              callback(service_id,
-                       *(down_cast<BluetoothSocket*>(
-                           multiplex_socket->GetVirtualSocket(service_id))));
-              callback_called = true;
+          if (multiplex_socket != nullptr) {
+            if (auto virtual_socket =
+                    multiplex_socket->GetVirtualSocket(service_id)) {
+              multiplex_sockets_.emplace(
+                  client_socket.GetRemoteDevice().GetAddress(),
+                  multiplex_socket);
+              MultiplexSocket::StopListeningForIncomingConnection(
+                  service_id, Medium::BLUETOOTH);
+              LOG(INFO) << "Multiplex virtaul socket created for "
+                        << client_socket.GetRemoteDevice().GetName();
+              if (callback) {
+                callback(service_id, *(down_cast<BluetoothSocket*>(
+                                         virtual_socket.get())));
+                callback_called = true;
+              }
             }
           }
         }
@@ -509,10 +508,11 @@ ErrorOr<BluetoothSocket> BluetoothClassic::Connect(
       if (it != multiplex_sockets_.end()) {
         MultiplexSocket* multiplex_socket = it->second;
         if (multiplex_socket->IsEnabled()) {
-          auto* virtual_socket =
+          std::shared_ptr<MediumSocket> virtual_socket =
               multiplex_socket->EstablishVirtualSocket(service_id);
           // Should not happen.
-          auto* bluetooth_socket = down_cast<BluetoothSocket*>(virtual_socket);
+          auto* bluetooth_socket =
+              down_cast<BluetoothSocket*>(virtual_socket.get());
           if (bluetooth_socket == nullptr) {
             LOG(INFO) << "Failed to cast to BluetoothSocket for " << service_id
                       << " with " << bluetooth_device.GetName();
@@ -607,9 +607,10 @@ ErrorOr<BluetoothSocket> BluetoothClassic::AttemptToConnect(
     MultiplexSocket* multiplex_socket = MultiplexSocket::CreateOutgoingSocket(
         std::move(physical_socket_ptr), service_id);
 
-    auto* virtual_socket = multiplex_socket->GetVirtualSocket(service_id);
-    // Should not happen.
-    auto* bluetooth_socket = down_cast<BluetoothSocket*>(virtual_socket);
+    std::shared_ptr<MediumSocket> virtual_socket =
+        multiplex_socket->GetVirtualSocket(service_id);
+
+    auto* bluetooth_socket = down_cast<BluetoothSocket*>(virtual_socket.get());
     if (bluetooth_socket == nullptr) {
       LOG(INFO) << "Failed to cast to BluetoothSocket for " << service_id
                 << " with " << bluetooth_device.GetName();
diff --git a/connections/implementation/mediums/bluetooth_classic.h b/connections/implementation/mediums/bluetooth_classic.h
index f16cf65f..a521d9bb 100644
--- a/connections/implementation/mediums/bluetooth_classic.h
+++ b/connections/implementation/mediums/bluetooth_classic.h
@@ -236,7 +236,13 @@ class BluetoothClassic {
       discovery_callbacks_ ABSL_GUARDED_BY(discovery_callbacks_mutex_);
 
   // Whether the multiplex feature is enabled.
-  bool is_multiplex_enabled_ = false;
+  bool is_multiplex_enabled_ =
+      NearbyFlags::GetInstance().GetBoolFlag(
+          config_package_nearby::nearby_connections_feature::
+              kEnableMultiplex) &&
+      NearbyFlags::GetInstance().GetBoolFlag(
+          config_package_nearby::nearby_connections_feature::
+              kEnableMultiplexBluetooth);
 
   // A map of Bluetooth MacAddress -> MultiplexSocket.
   absl::flat_hash_map<MacAddress, mediums::multiplex::MultiplexSocket*>
diff --git a/connections/implementation/mediums/multiplex/multiplex_socket.cc b/connections/implementation/mediums/multiplex/multiplex_socket.cc
index befe36e9..888efb82 100644
--- a/connections/implementation/mediums/multiplex/multiplex_socket.cc
+++ b/connections/implementation/mediums/multiplex/multiplex_socket.cc
@@ -16,7 +16,6 @@
 
 #include <cstdint>
 #include <memory>
-#include <new>
 #include <string>
 #include <type_traits>
 #include <utility>
@@ -215,7 +214,7 @@ MultiplexSocket* MultiplexSocket::CreateOutgoingSocket(
                               Utils::GenerateSalt());
 }
 
-MediumSocket* MultiplexSocket::CreateFirstVirtualSocket(
+std::shared_ptr<MediumSocket> MultiplexSocket::CreateFirstVirtualSocket(
     const std::string& service_id, const std::string& service_id_hash_salt) {
   auto output_stream =
       multiplex_output_stream_.CreateVirtualOutputStreamForFirstVirtualSocket(
@@ -227,9 +226,14 @@ MediumSocket* MultiplexSocket::CreateFirstVirtualSocket(
   LOG(INFO) << __func__ << " for service_id=" << service_id
             << ", salt=" << service_id_hash_salt
             << ", salted_service_id_hash_key=" << salted_service_id_hash_key;
-  MediumSocket* virtual_socket = physical_socket_ptr_->CreateVirtualSocket(
+  MediumSocket* virtual_socket_ptr = physical_socket_ptr_->CreateVirtualSocket(
       salted_service_id_hash_key, output_stream, medium_, &virtual_sockets_);
 
+  if (virtual_socket_ptr == nullptr) {
+    return nullptr;
+  }
+  std::shared_ptr<MediumSocket> virtual_socket =
+      virtual_sockets_[salted_service_id_hash_key];
   virtual_socket->AddOnSocketClosedListener(
       std::make_unique<absl::AnyInvocable<void()>>(
           [this, service_id]() { OnVirtualSocketClosed(service_id); }));
@@ -242,7 +246,7 @@ MediumSocket* MultiplexSocket::CreateFirstVirtualSocket(
   return virtual_socket;
 }
 
-MediumSocket* MultiplexSocket::CreateVirtualSocket(
+std::shared_ptr<MediumSocket> MultiplexSocket::CreateVirtualSocket(
     const std::string& service_id, const std::string& service_id_hash_salt) {
   auto output_stream = multiplex_output_stream_.CreateVirtualOutputStream(
       service_id, service_id_hash_salt);
@@ -254,9 +258,14 @@ MediumSocket* MultiplexSocket::CreateVirtualSocket(
             << ", salt=" << service_id_hash_salt
             << ", salted_service_id_hash_key=" << salted_service_id_hash_key;
 
-  MediumSocket* virtual_socket = physical_socket_ptr_->CreateVirtualSocket(
+  MediumSocket* virtual_socket_ptr = physical_socket_ptr_->CreateVirtualSocket(
       salted_service_id_hash_key, output_stream, medium_, &virtual_sockets_);
 
+  if (virtual_socket_ptr == nullptr) {
+    return nullptr;
+  }
+  std::shared_ptr<MediumSocket> virtual_socket =
+      virtual_sockets_[salted_service_id_hash_key];
   virtual_socket->AddOnSocketClosedListener(
       std::make_unique<absl::AnyInvocable<void()>>(
           [this, service_id]() { OnVirtualSocketClosed(service_id); }));
@@ -264,7 +273,8 @@ MediumSocket* MultiplexSocket::CreateVirtualSocket(
   return virtual_socket;
 }
 
-MediumSocket* MultiplexSocket::GetVirtualSocket(const std::string& service_id) {
+std::shared_ptr<MediumSocket> MultiplexSocket::GetVirtualSocket(
+    const std::string& service_id) {
   MutexLock lock(&virtual_socket_mutex_);
   LOG(INFO) << __func__ << " service_id=" << service_id << ", Salt="
             << multiplex_output_stream_.GetServiceIdHashSalt(service_id)
@@ -275,7 +285,7 @@ MediumSocket* MultiplexSocket::GetVirtualSocket(const std::string& service_id) {
     LOG(INFO) << "Not found!";
     return nullptr;
   }
-  return item->second.get();
+  return item->second;
 }
 
 int MultiplexSocket::GetVirtualSocketCount() {
@@ -305,7 +315,7 @@ void MultiplexSocket::UnRegisterConnectionResponse(
   connection_response_futures_.erase(service_id);
 }
 
-MediumSocket* MultiplexSocket::EstablishVirtualSocket(
+std::shared_ptr<MediumSocket> MultiplexSocket::EstablishVirtualSocket(
     const std::string& service_id) {
   if (!IsEnabled()) {
     LOG(ERROR)
@@ -555,7 +565,7 @@ void MultiplexSocket::HandleConnectionRequest(
       << "EstablishVirtualSocket after local device accept the connection "
          "with serviceId="
       << listening_service_id;
-  MediumSocket* virtual_socket =
+  std::shared_ptr<MediumSocket> virtual_socket =
       CreateVirtualSocket(listening_service_id, service_id_hash_salt);
   (*incoming_connection_callback)(std::move(listening_service_id),
                                   virtual_socket);
@@ -620,13 +630,13 @@ void MultiplexSocket::HandleDataFrame(const ByteArray& salted_service_id_hash,
                                       const MultiplexDataFrame& frame) {
   std::string salted_service_id_hash_key =
       GenerateServiceIdHashKey(salted_service_id_hash);
-  MediumSocket* virtual_socket = nullptr;
+  std::shared_ptr<MediumSocket> virtual_socket = nullptr;
   if (service_id_hash_salt.empty()) {
     {
       MutexLock lock(&virtual_socket_mutex_);
       auto item = virtual_sockets_.find(salted_service_id_hash_key);
       if (item != virtual_sockets_.end()) {
-        virtual_socket = item->second.get();
+        virtual_socket = item->second;
       }
     }
   } else {
@@ -659,7 +669,8 @@ void MultiplexSocket::OnVirtualSocketClosed(const std::string& service_id) {
   RunOffloadThread(
       "VirtualSocketClosed", [this, service_id, &latch, &shutdown]() {
         LOG(INFO) << "Try to close Virtual socket: " << service_id;
-        MediumSocket* virtual_socket = GetVirtualSocket(service_id);
+        std::shared_ptr<MediumSocket> virtual_socket =
+            GetVirtualSocket(service_id);
         {
           MutexLock lock(&virtual_socket_mutex_);
           LOG(INFO) << "virtual_socket:" << virtual_socket;
@@ -700,7 +711,7 @@ void MultiplexSocket::OnVirtualSocketClosed(const std::string& service_id) {
   }
 }
 
-MediumSocket* MultiplexSocket::ReMapAndGetVirtualSocket(
+std::shared_ptr<MediumSocket> MultiplexSocket::ReMapAndGetVirtualSocket(
     const ByteArray& salted_service_id_hash,
     const std::string& service_id_hash_salt) {
   std::string salted_service_id_hash_key =
@@ -722,7 +733,7 @@ MediumSocket* MultiplexSocket::ReMapAndGetVirtualSocket(
       }
       if ((service_id_hash_salt == kFakeSalt) ||
           (hash_key == salted_service_id_hash_key)) {
-        return virtual_socket.get();
+        return virtual_socket;
       } else {
         LOG(INFO) << "Remap the virtualSockets.";
         output_stream->SetserviceIdHashSalt(service_id_hash_salt);
@@ -731,7 +742,7 @@ MediumSocket* MultiplexSocket::ReMapAndGetVirtualSocket(
         virtual_sockets_.erase(hash_key);
         virtual_sockets_[salted_service_id_hash_key] = virtual_socket_tmp;
         ListVirtualSocket();
-        return virtual_socket_tmp.get();
+        return virtual_socket_tmp;
       }
     }
   }
diff --git a/connections/implementation/mediums/multiplex/multiplex_socket.h b/connections/implementation/mediums/multiplex/multiplex_socket.h
index 5e65c5d4..ed3fdd7d 100644
--- a/connections/implementation/mediums/multiplex/multiplex_socket.h
+++ b/connections/implementation/mediums/multiplex/multiplex_socket.h
@@ -44,7 +44,7 @@ namespace multiplex {
 
 using MultiplexEnbaleCb = absl::AnyInvocable<void()>;
 using MultiplexIncomingConnectionCb = absl::AnyInvocable<void(
-    const std::string& service_id, MediumSocket* socket)>;
+    const std::string& service_id, std::shared_ptr<MediumSocket> socket)>;
 
 class MultiplexSocket {
  public:
@@ -82,7 +82,7 @@ class MultiplexSocket {
       const std::string& service_id,
       ::location::nearby::proto::connections::Medium type,
       absl::AnyInvocable<void(const std::string& service_id,
-                              MediumSocket* socket)>
+                              std::shared_ptr<MediumSocket> socket)>
           incoming_connection_cb);
 
   // Stops listening for incoming multiplex connection for {@code service_id} on
@@ -98,7 +98,7 @@ class MultiplexSocket {
   }
 
   // Gets the virtual socket by service id.
-  MediumSocket* GetVirtualSocket(const std::string& service_id);
+  std::shared_ptr<MediumSocket> GetVirtualSocket(const std::string& service_id);
   // Gets the virtual socket count.
   int GetVirtualSocketCount();
 
@@ -106,7 +106,8 @@ class MultiplexSocket {
       ABSL_EXCLUSIVE_LOCKS_REQUIRED(virtual_socket_mutex_);
 
   // Establishes the virtual socket by service id.
-  MediumSocket* EstablishVirtualSocket(const std::string& service_id);
+  std::shared_ptr<MediumSocket> EstablishVirtualSocket(
+      const std::string& service_id);
   // Shuts down the multiplex socket.
   void Shutdown();
   bool IsShutdown() { return is_shutdown_; }
@@ -118,11 +119,11 @@ class MultiplexSocket {
 
   // Creates the first virtual socket for the service id. The first virtual
   // socket is created by the sender.
-  MediumSocket* CreateFirstVirtualSocket(
+  std::shared_ptr<MediumSocket> CreateFirstVirtualSocket(
       const std::string& service_id, const std::string& service_id_hash_salt);
   // Creates the virtual socket for the service id.
-  MediumSocket* CreateVirtualSocket(const std::string& service_id,
-                                    const std::string& service_id_hash_salt);
+  std::shared_ptr<MediumSocket> CreateVirtualSocket(
+      const std::string& service_id, const std::string& service_id_hash_salt);
   // Registers the connection response future for the service id.
   std::shared_ptr<Future<::location::nearby::mediums::ConnectionResponseFrame::
                              ConnectionResponseCode>>
@@ -157,7 +158,7 @@ class MultiplexSocket {
   // Handles the physical socket closed.
   void OnPhysicalSocketClosed();
   // Remaps and gets the virtual socket by service id hash.
-  MediumSocket* ReMapAndGetVirtualSocket(
+  std::shared_ptr<MediumSocket> ReMapAndGetVirtualSocket(
       const ByteArray& salted_service_id_hash,
       const std::string& service_id_hash_salt);
   // Handles the virtual socket closed.
diff --git a/connections/implementation/mediums/multiplex/multiplex_socket_test.cc b/connections/implementation/mediums/multiplex/multiplex_socket_test.cc
index aba53de0..cf382eef 100644
--- a/connections/implementation/mediums/multiplex/multiplex_socket_test.cc
+++ b/connections/implementation/mediums/multiplex/multiplex_socket_test.cc
@@ -38,6 +38,7 @@
 #include "internal/platform/pipe.h"
 #include "internal/platform/single_thread_executor.h"
 #include "internal/platform/socket.h"
+#include "internal/platform/types.h"
 #include "proto/connections_enums.proto.h"
 
 namespace nearby {
@@ -89,9 +90,9 @@ class FakeSocket : public MediumSocket {
 
   InputStream& GetInputStream() override { return *reader_1_; }
   OutputStream& GetOutputStream() override {
-    return IsVirtualSocket() ? *virtual_output_stream_
-                             : *writer_2_;
-  }  Exception Close() override {
+    return IsVirtualSocket() ? *virtual_output_stream_ : *writer_2_;
+  }
+  Exception Close() override {
     if (IsVirtualSocket()) {
       LOG(INFO) << "Multiplex: Closing virtual socket: " << this;
       CloseLocal();
@@ -117,8 +118,8 @@ class FakeSocket : public MediumSocket {
     }
 
     auto virtual_socket = std::make_shared<FakeSocket>(medium, outputstream);
-    LOG(WARNING) << "Created the virtual socket for Medium: "
-                 << Medium_Name(virtual_socket->GetMedium());
+    LOG(INFO) << "Created the virtual socket for Medium: "
+              << Medium_Name(virtual_socket->GetMedium());
 
     if (virtual_sockets_ptr_ == nullptr) {
       virtual_sockets_ptr_ = virtual_sockets_ptr;
@@ -163,12 +164,12 @@ TEST(MultiplexSocketTest, CreateIncomingSocketSuccess) {
                                                       Medium::BLUETOOTH);
   MultiplexSocket::ListenForIncomingConnection(
       std::string(SERVICE_ID_1), Medium::BLUETOOTH,
-      [](const std::string& service_id, MediumSocket* socket) {
+      [](const std::string& service_id, std::shared_ptr<MediumSocket> socket) {
         LOG(INFO) << "Incoming connection for service_id: " << service_id;
       });
   MultiplexSocket::ListenForIncomingConnection(
       std::string(SERVICE_ID_2), Medium::BLUETOOTH,
-      [](const std::string& service_id, MediumSocket* socket) {
+      [](const std::string& service_id, std::shared_ptr<MediumSocket> socket) {
         LOG(INFO) << "Incoming connection for service_id: " << service_id;
       });
 
@@ -181,13 +182,11 @@ TEST(MultiplexSocketTest, CreateIncomingSocketSuccess) {
           fake_socket_ptr, std::string(SERVICE_ID_2), /*first_frame_len*/ 0);
   ASSERT_EQ(multiplex_socket_incoming_2, multiplex_socket_incoming);
 
+  std::shared_ptr<MediumSocket> virtual_socket_shared =
+      multiplex_socket_incoming->GetVirtualSocket(std::string(SERVICE_ID_1));
+  ASSERT_NE(virtual_socket_shared, nullptr);
   FakeSocket* virtual_socket =
-      (FakeSocket*)multiplex_socket_incoming->GetVirtualSocket(
-          std::string(SERVICE_ID_1));
-  if (virtual_socket == nullptr) {
-    LOG(INFO) << "Virtual socket not found for " << SERVICE_ID_1;
-    return;
-  }
+      down_cast<FakeSocket*>(virtual_socket_shared.get());
 
   SingleThreadExecutor executor;
   FakeSocket* socket = fake_socket_ptr.get();
@@ -240,12 +239,12 @@ TEST(MultiplexSocketTest, CreateIncomingVirtualSocketSuccess) {
                                                       Medium::WIFI_LAN);
   MultiplexSocket::ListenForIncomingConnection(
       std::string(SERVICE_ID_1), Medium::WIFI_LAN,
-      [](const std::string& service_id, MediumSocket* socket) {
+      [](const std::string& service_id, std::shared_ptr<MediumSocket> socket) {
         LOG(INFO) << "Incoming connection for service_id: " << service_id;
       });
   MultiplexSocket::ListenForIncomingConnection(
       std::string(SERVICE_ID_2), Medium::WIFI_LAN,
-      [](const std::string& service_id, MediumSocket* socket) {
+      [](const std::string& service_id, std::shared_ptr<MediumSocket> socket) {
         LOG(INFO) << "Incoming connection for service_id: " << service_id;
       });
 
@@ -254,13 +253,11 @@ TEST(MultiplexSocketTest, CreateIncomingVirtualSocketSuccess) {
           fake_socket_ptr, std::string(SERVICE_ID_1), /*first_frame_len*/ 0);
   ASSERT_NE(multiplex_socket_incoming, nullptr);
 
+  std::shared_ptr<MediumSocket> virtual_socket_shared =
+      multiplex_socket_incoming->GetVirtualSocket(std::string(SERVICE_ID_1));
+  ASSERT_NE(virtual_socket_shared, nullptr);
   FakeSocket* virtual_socket =
-      (FakeSocket*)multiplex_socket_incoming->GetVirtualSocket(
-          std::string(SERVICE_ID_1));
-  if (virtual_socket == nullptr) {
-    LOG(INFO) << "Virtual socket not found for " << SERVICE_ID_1;
-    return;
-  }
+      down_cast<FakeSocket*>(virtual_socket_shared.get());
 
   SingleThreadExecutor executor;
   FakeSocket* socket = fake_socket_ptr.get();
@@ -296,42 +293,57 @@ TEST(MultiplexSocketTest,
       fake_socket_ptr, std::string(SERVICE_ID_2));
   ASSERT_EQ(multiplex_socket_2, multiplex_socket);
   multiplex_socket->Enable();
-  FakeSocket* virtual_socket = (FakeSocket*)multiplex_socket->GetVirtualSocket(
-      std::string(SERVICE_ID_1));
-  if (virtual_socket == nullptr) {
-    LOG(INFO) << "Virtual socket not found for " << SERVICE_ID_1;
-    return;
-  }
+  std::shared_ptr<MediumSocket> virtual_socket_shared =
+      multiplex_socket->GetVirtualSocket(std::string(SERVICE_ID_1));
+  ASSERT_NE(virtual_socket_shared, nullptr);
+  FakeSocket* virtual_socket =
+      down_cast<FakeSocket*>(virtual_socket_shared.get());
 
+  // This is a timeout test, the real timeout is 3s which is too long for a
+  // unit test, so we set a short timeout for flakiness test to avoid long wait
+  // time.
+  auto flags = FeatureFlags::GetInstance().GetFlags();
+  auto original_flags = flags;
+  flags.multiplex_socket_connection_response_timeout_millis =
+      absl::Milliseconds(200);
+  FeatureFlags::GetMutableInstanceForTesting().SetFlags(flags);
+
+  CountDownLatch latch(2);
   SingleThreadExecutor establish_socket_executor;
-  establish_socket_executor.Execute([&multiplex_socket]() {
+  establish_socket_executor.Execute([&multiplex_socket, &latch]() {
     LOG(INFO) << "EstablishVirtualSocket";
-    MediumSocket* socket =
+    std::shared_ptr<MediumSocket> socket =
         multiplex_socket->EstablishVirtualSocket(std::string(SERVICE_ID_2));
     LOG(INFO) << "EstablishVirtualSocket finished";
     EXPECT_EQ(socket, nullptr);
+    latch.CountDown();
   });
 
   SingleThreadExecutor read_executor;
-  read_executor.Execute([&multiplex_socket, &fake_socket_ptr]() {
+  read_executor.Execute([&multiplex_socket, &fake_socket_ptr, &latch]() {
     auto reader = fake_socket_ptr->reader_2_.get();
     LOG(INFO) << "reader_2_ Read start";
     ExceptionOr<std::int32_t> read_int = Base64Utils::ReadInt(reader);
     if (!read_int.ok()) {
       ADD_FAILURE() << "Failed to read. Exception:" << read_int.exception();
+    } else {
+      auto length = read_int.result();
+      LOG(INFO) << " length:" << length;
+      EXPECT_GT(length, 0);
     }
-    auto length = read_int.result();
-    LOG(INFO) << " length:" << length;
-    EXPECT_GT(length, 0);
     EXPECT_EQ(multiplex_socket->GetVirtualSocket(std::string(SERVICE_ID_2)),
               nullptr);
+    latch.CountDown();
   });
 
-  absl::SleepFor(absl::Milliseconds(300));
+  EXPECT_TRUE(latch.Await(absl::Seconds(1)).result());
   EXPECT_EQ(multiplex_socket->GetVirtualSocketCount(), 1);
   virtual_socket->Close();
   EXPECT_EQ(multiplex_socket->GetVirtualSocketCount(), 0);
   multiplex_socket->ShutdownAll();
+
+  // Restore the original flags.
+  FeatureFlags::GetMutableInstanceForTesting().SetFlags(original_flags);
 }
 
 TEST(MultiplexSocketTest, EstablishVirtualSocket_RemoteAccepted) {
@@ -352,7 +364,7 @@ TEST(MultiplexSocketTest, EstablishVirtualSocket_RemoteAccepted) {
   CountDownLatch latch(1);
   executor.Execute([&multiplex_socket, &latch]() {
     LOG(INFO) << "EstablishVirtualSocket";
-    MediumSocket* socket =
+    std::shared_ptr<MediumSocket> socket =
         multiplex_socket->EstablishVirtualSocket(std::string(SERVICE_ID_2));
     EXPECT_EQ(socket, nullptr);
     latch.CountDown();
@@ -362,7 +374,7 @@ TEST(MultiplexSocketTest, EstablishVirtualSocket_RemoteAccepted) {
   multiplex_socket->Enable();
   executor.Execute([&multiplex_socket]() {
     LOG(INFO) << "EstablishVirtualSocket";
-    MediumSocket* socket =
+    std::shared_ptr<MediumSocket> socket =
         multiplex_socket->EstablishVirtualSocket(std::string(SERVICE_ID_2));
     EXPECT_NE(socket, nullptr);
   });
diff --git a/connections/implementation/mediums/wifi_lan.cc b/connections/implementation/mediums/wifi_lan.cc
index ff3c58ad..9f30ea5d 100644
--- a/connections/implementation/mediums/wifi_lan.cc
+++ b/connections/implementation/mediums/wifi_lan.cc
@@ -270,10 +270,10 @@ ErrorOr<int> WifiLan::StartAcceptingConnectionsLocked(
     MultiplexSocket::ListenForIncomingConnection(
         service_id, Medium::WIFI_LAN,
         [&callback](const std::string& listening_service_id,
-                    MediumSocket* virtual_socket) mutable {
+                    std::shared_ptr<MediumSocket> virtual_socket) mutable {
           if (callback) {
             callback(listening_service_id,
-                     *(down_cast<WifiLanSocket*>(virtual_socket)));
+                     *(down_cast<WifiLanSocket*>(virtual_socket.get())));
           }
         });
   }
@@ -322,20 +322,21 @@ ErrorOr<int> WifiLan::StartAcceptingConnectionsLocked(
               MultiplexSocket* multiplex_socket =
                   MultiplexSocket::CreateIncomingSocket(
                       physical_socket_ptr, service_id, read_int.result());
-              if (multiplex_socket != nullptr &&
-                  multiplex_socket->GetVirtualSocket(service_id)) {
-                multiplex_sockets_.emplace(server_socket.GetIPAddress(),
-                                           multiplex_socket);
-                MultiplexSocket::StopListeningForIncomingConnection(
-                    service_id, Medium::WIFI_LAN);
-                LOG(INFO) << "Multiplex virtaul socket created for "
-                          << server_socket.GetIPAddress();
-                if (callback) {
-                  callback(
-                      service_id,
-                      *(down_cast<WifiLanSocket*>(
-                          multiplex_socket->GetVirtualSocket(service_id))));
-                  callback_called = true;
+              if (multiplex_socket != nullptr) {
+                std::shared_ptr<MediumSocket> virtual_socket =
+                    multiplex_socket->GetVirtualSocket(service_id);
+                if (virtual_socket) {
+                  multiplex_sockets_.emplace(server_socket.GetIPAddress(),
+                                             multiplex_socket);
+                  MultiplexSocket::StopListeningForIncomingConnection(
+                      service_id, Medium::WIFI_LAN);
+                  LOG(INFO) << "Multiplex virtaul socket created for "
+                            << server_socket.GetIPAddress();
+                  if (callback) {
+                    callback(service_id, *(down_cast<WifiLanSocket*>(
+                                             virtual_socket.get())));
+                    callback_called = true;
+                  }
                 }
               }
             }
@@ -568,10 +569,10 @@ ExceptionOr<WifiLanSocket> WifiLan::ConnectWithMultiplexSocketLocked(
         return ExceptionOr<WifiLanSocket>(Exception::kFailed);
       }
       if (multiplex_socket->IsEnabled()) {
-        auto* virtual_socket =
+        std::shared_ptr<MediumSocket> virtual_socket =
             multiplex_socket->EstablishVirtualSocket(service_id);
         // Should not happen.
-        auto* wlan_socket = down_cast<WifiLanSocket*>(virtual_socket);
+        auto* wlan_socket = down_cast<WifiLanSocket*>(virtual_socket.get());
         if (wlan_socket == nullptr) {
           LOG(INFO) << "Failed to cast to WifiLanSocket for " << service_id
                     << " with ip_address: "
@@ -595,9 +596,9 @@ ExceptionOr<WifiLanSocket> WifiLan::CreateOutgoingMultiplexSocketLocked(
     MultiplexSocket* multiplex_socket =
         MultiplexSocket::CreateOutgoingSocket(physical_socket_ptr, service_id);
 
-    auto* virtual_socket = multiplex_socket->GetVirtualSocket(service_id);
-    // Should not happen.
-    auto* wlan_socket = down_cast<WifiLanSocket*>(virtual_socket);
+    std::shared_ptr<MediumSocket> virtual_socket =
+        multiplex_socket->GetVirtualSocket(service_id);
+    auto* wlan_socket = down_cast<WifiLanSocket*>(virtual_socket.get());
     if (wlan_socket == nullptr) {
       LOG(INFO) << "Failed to cast to WifiLanSocket for " << service_id
                 << " with ip_address: "
diff --git a/connections/implementation/mediums/wifi_lan.h b/connections/implementation/mediums/wifi_lan.h
index 3637b158..4dffc2c1 100644
--- a/connections/implementation/mediums/wifi_lan.h
+++ b/connections/implementation/mediums/wifi_lan.h
@@ -219,7 +219,13 @@ class WifiLan {
       ABSL_GUARDED_BY(mutex_);
 
   // Whether the multiplex feature is enabled.
-  bool is_multiplex_enabled_ = false;
+  bool is_multiplex_enabled_ =
+      NearbyFlags::GetInstance().GetBoolFlag(
+          config_package_nearby::nearby_connections_feature::
+              kEnableMultiplex) &&
+      NearbyFlags::GetInstance().GetBoolFlag(
+          config_package_nearby::nearby_connections_feature::
+              kEnableMultiplexWifiLan);
 
   // A map of IpAddress -> MultiplexSocket.
   absl::flat_hash_map<std::string, mediums::multiplex::MultiplexSocket*>