diff --git a/connections/implementation/BUILD b/connections/implementation/BUILD
index c3723458..1dd42abe 100644
--- a/connections/implementation/BUILD
+++ b/connections/implementation/BUILD
@@ -251,6 +251,7 @@ cc_library(
         "//internal/interop:authentication_transport_interface",
         "//internal/interop:device",
         "//internal/platform:base",
+        "//internal/platform:cancellation_flag",
         "//internal/platform:comm",
         "//internal/platform:connection_info",
         "//internal/platform:logging",
@@ -478,6 +479,7 @@ cc_test(
         "//internal/platform/implementation/g3",  # build_cleaner: keep
         "//proto:connections_enums_cc_proto",
         "@com_github_protobuf_matchers//protobuf-matchers",
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/time",
diff --git a/connections/implementation/client_proxy.cc b/connections/implementation/client_proxy.cc
index e4737d0f..31bfc12b 100644
--- a/connections/implementation/client_proxy.cc
+++ b/connections/implementation/client_proxy.cc
@@ -190,6 +190,7 @@ const NearbyDevice* ClientProxy::GetLocalDevice() {
 }
 
 std::string ClientProxy::GetConnectionToken(const std::string& endpoint_id) {
+  MutexLock lock(&mutex_);
   ConnectionPair* item = LookupConnection(endpoint_id);
   if (item != nullptr) {
     return item->first.connection_token;
@@ -220,6 +221,7 @@ std::string ClientProxy::GetSavePath(
 
 std::optional<MacAddress> ClientProxy::GetBluetoothMacAddress(
     const std::string& endpoint_id) {
+  MutexLock lock(&mutex_);
   auto item = bluetooth_mac_addresses_.find(endpoint_id);
   if (item != bluetooth_mac_addresses_.end()) return item->second;
   return std::nullopt;
@@ -227,6 +229,7 @@ std::optional<MacAddress> ClientProxy::GetBluetoothMacAddress(
 
 void ClientProxy::SetBluetoothMacAddress(const std::string& endpoint_id,
                                          MacAddress bluetooth_mac_address) {
+  MutexLock lock(&mutex_);
   bluetooth_mac_addresses_[endpoint_id] = bluetooth_mac_address;
 }
 
@@ -878,16 +881,19 @@ bool ClientProxy::IsConnectionRejected(const std::string& endpoint_id) const {
 }
 
 bool ClientProxy::LocalConnectionIsAccepted(std::string endpoint_id) const {
+  MutexLock lock(&mutex_);
   return ConnectionStatusesContains(
       endpoint_id, ClientProxy::Connection::kLocalEndpointAccepted);
 }
 
 bool ClientProxy::RemoteConnectionIsAccepted(std::string endpoint_id) const {
+  MutexLock lock(&mutex_);
   return ConnectionStatusesContains(
       endpoint_id, ClientProxy::Connection::kRemoteEndpointAccepted);
 }
 
 bool ClientProxy::AutoUpgradeBandwidth() const {
+  MutexLock lock(&mutex_);
   bool result = false;
   if (IsAdvertising() && (GetAdvertisingOptions().strategy.IsNone() ||
                           GetAdvertisingOptions().auto_upgrade_bandwidth)) {
@@ -902,6 +908,7 @@ bool ClientProxy::AutoUpgradeBandwidth() const {
 }
 
 bool ClientProxy::ShouldEnforceTopologyConstraints() const {
+  MutexLock lock(&mutex_);
   bool result = false;
   if (IsAdvertising() &&
       (GetAdvertisingOptions().strategy.IsNone() ||
@@ -922,6 +929,7 @@ void ClientProxy::AddCancellationFlag(const std::string& endpoint_id) {
     return;
   }
 
+  MutexLock lock(&mutex_);
   auto item = cancellation_flags_.find(endpoint_id);
   if (item != cancellation_flags_.end()) {
     // A new flag may be added to the map with the same endpoint, even if a
@@ -936,19 +944,21 @@ void ClientProxy::AddCancellationFlag(const std::string& endpoint_id) {
     return;
   }
   cancellation_flags_.emplace(endpoint_id,
-                              std::make_unique<CancellationFlag>());
+                              std::make_shared<CancellationFlag>());
 }
 
-CancellationFlag* ClientProxy::GetCancellationFlag(
+std::shared_ptr<CancellationFlag> ClientProxy::GetCancellationFlag(
     const std::string& endpoint_id) {
+  MutexLock lock(&mutex_);
   const auto item = cancellation_flags_.find(endpoint_id);
   if (item == cancellation_flags_.end()) {
-    return default_cancellation_flag_.get();
+    return default_cancellation_flag_;
   }
-  return item->second.get();
+  return item->second;
 }
 
 void ClientProxy::CancelEndpoint(const std::string& endpoint_id) {
+  MutexLock lock(&mutex_);
   const auto item = cancellation_flags_.find(endpoint_id);
   if (item != cancellation_flags_.end()) {
     item->second->Cancel();
@@ -959,6 +969,7 @@ const OsInfo& ClientProxy::GetLocalOsInfo() const { return local_os_info_; }
 
 std::optional<OsInfo> ClientProxy::GetRemoteOsInfo(
     absl::string_view endpoint_id) const {
+  MutexLock lock(&mutex_);
   const ConnectionPair* item = LookupConnection(endpoint_id);
   if (item != nullptr) {
     return item->first.os_info;
@@ -968,11 +979,13 @@ std::optional<OsInfo> ClientProxy::GetRemoteOsInfo(
 
 void ClientProxy::SetLocalOsType(
     const location::nearby::connections::OsInfo::OsType& os_type) {
+  MutexLock lock(&mutex_);
   local_os_info_.set_type(os_type);
 }
 
 void ClientProxy::SetRemoteOsInfo(absl::string_view endpoint_id,
                                   const OsInfo& remote_os_info) {
+  MutexLock lock(&mutex_);
   ConnectionPair* item = LookupConnection(endpoint_id);
   if (item != nullptr) {
     item->first.os_info.emplace(remote_os_info);
@@ -1019,8 +1032,9 @@ bool ClientProxy::IsPayloadReceivedAckEnabled(absl::string_view endpoint_id) {
 }
 
 void ClientProxy::CancelAllEndpoints() {
+  MutexLock lock(&mutex_);
   for (const auto& item : cancellation_flags_) {
-    CancellationFlag* cancellation_flag = item.second.get();
+    std::shared_ptr<CancellationFlag> cancellation_flag = item.second;
     if (cancellation_flag->Cancelled()) {
       continue;
     }
@@ -1123,14 +1137,17 @@ void ClientProxy::AppendConnectionStatus(const std::string& endpoint_id,
 }
 
 AdvertisingOptions ClientProxy::GetAdvertisingOptions() const {
+  MutexLock lock(&mutex_);
   return advertising_options_;
 }
 
 DiscoveryOptions ClientProxy::GetDiscoveryOptions() const {
+  MutexLock lock(&mutex_);
   return discovery_options_;
 }
 
 v3::ConnectionListeningOptions ClientProxy::GetListeningOptions() const {
+  MutexLock lock(&mutex_);
   return listening_options_;
 }
 
@@ -1208,11 +1225,13 @@ OsInfo::OsType ClientProxy::OSNameToOsInfoType(api::OSName osName) {
 }
 
 std::int32_t ClientProxy::GetLocalMultiplexSocketBitmask() const {
+  MutexLock lock(&mutex_);
   return 0;
 }
 
 void ClientProxy::SetRemoteMultiplexSocketBitmask(
     absl::string_view endpoint_id, int remote_multiplex_socket_bitmask) {
+  MutexLock lock(&mutex_);
   ConnectionPair* item = LookupConnection(endpoint_id);
   if (item != nullptr) {
     item->first.remote_multiplex_socket_bitmask =
@@ -1223,6 +1242,7 @@ void ClientProxy::SetRemoteMultiplexSocketBitmask(
 }
 
 bool ClientProxy::IsLocalMultiplexSocketSupported(Medium medium) {
+  MutexLock lock(&mutex_);
   int bitmask = GetLocalMultiplexSocketBitmask();
   switch (medium) {
     case Medium::BLUETOOTH:
@@ -1238,6 +1258,7 @@ bool ClientProxy::IsLocalMultiplexSocketSupported(Medium medium) {
 
 std::optional<std::int32_t> ClientProxy::GetRemoteMultiplexSocketBitmask(
     absl::string_view endpoint_id) const {
+  MutexLock lock(&mutex_);
   const ConnectionPair* item = LookupConnection(endpoint_id);
   if (item != nullptr) {
     return item->first.remote_multiplex_socket_bitmask;
@@ -1247,6 +1268,7 @@ std::optional<std::int32_t> ClientProxy::GetRemoteMultiplexSocketBitmask(
 
 bool ClientProxy::IsMultiplexSocketSupported(absl::string_view endpoint_id,
                                              Medium medium) {
+  MutexLock lock(&mutex_);
   ConnectionPair* item = LookupConnection(endpoint_id);
   if (item == nullptr) {
     return false;
@@ -1264,20 +1286,31 @@ bool ClientProxy::IsMultiplexSocketSupported(absl::string_view endpoint_id,
   }
 }
 
-bool ClientProxy::GetWebRtcNonCellular() { return webrtc_non_cellular_; }
+bool ClientProxy::GetWebRtcNonCellular() {
+  MutexLock lock(&mutex_);
+  return webrtc_non_cellular_;
+}
 
 void ClientProxy::SetWebRtcNonCellular(bool webrtc_non_cellular) {
+  MutexLock lock(&mutex_);
   VLOG(1) << "ClientProxy: client=" << GetClientId()
           << (webrtc_non_cellular ? " disallow" : " allow")
           << " to use mobile data.";
   webrtc_non_cellular_ = webrtc_non_cellular;
 }
 
-bool ClientProxy::IsDctEnabled() const { return is_dct_enabled_; }
+bool ClientProxy::IsDctEnabled() const {
+  MutexLock lock(&mutex_);
+  return is_dct_enabled_;
+}
 
-uint8_t ClientProxy::GetDctDedup() const { return dct_dedup_; }
+uint8_t ClientProxy::GetDctDedup() const {
+  MutexLock lock(&mutex_);
+  return dct_dedup_;
+}
 
 void ClientProxy::UpdateDctDeviceName(absl::string_view device_name) {
+  MutexLock lock(&mutex_);
   if (!dct_device_name_.empty() && dct_device_name_ != device_name) {
     // Need to update dedup value if device name is changed.
     absl::BitGen bitgen;
@@ -1299,6 +1332,7 @@ void ClientProxy::UpdateDctDeviceName(absl::string_view device_name) {
 
 std::optional<MediumRole> ClientProxy::GetMediumRole(
     absl::string_view endpoint_id) const {
+  MutexLock lock(&mutex_);
   const ConnectionPair* item = LookupConnection(endpoint_id);
   if (item != nullptr) {
     return item->first.connection_options.connection_info.medium_role;
@@ -1307,6 +1341,7 @@ std::optional<MediumRole> ClientProxy::GetMediumRole(
 }
 
 std::optional<std::string> ClientProxy::GetEndpointIdForDct() const {
+  MutexLock lock(&mutex_);
   if (dct_endpoint_id_.empty()) {
     return std::nullopt;
   }
diff --git a/connections/implementation/client_proxy.h b/connections/implementation/client_proxy.h
index e4929736..e7010739 100644
--- a/connections/implementation/client_proxy.h
+++ b/connections/implementation/client_proxy.h
@@ -262,7 +262,8 @@ class ClientProxy final {
   // Adds a CancellationFlag for endpoint id.
   void AddCancellationFlag(const std::string& endpoint_id);
   // Returns the CancellationFlag for endpoint id,
-  CancellationFlag* GetCancellationFlag(const std::string& endpoint_id);
+  std::shared_ptr<CancellationFlag> GetCancellationFlag(
+      const std::string& endpoint_id);
   // Sets the CancellationFlag to true for endpoint id.
   void CancelEndpoint(const std::string& endpoint_id);
   // Cancels all CancellationFlags.
@@ -517,11 +518,11 @@ class ClientProxy final {
   // Maps endpoint_id to CancellationFlag. CancellationFlags are passed around
   // as raw pointers to other classes in Nearby Connections, so it is important
   // that objects in this map are not cleared, even if they are cancelled.
-  absl::flat_hash_map<std::string, std::unique_ptr<CancellationFlag>>
+  absl::flat_hash_map<std::string, std::shared_ptr<CancellationFlag>>
       cancellation_flags_;
   // A default cancellation flag with isCancelled set be true.
-  std::unique_ptr<CancellationFlag> default_cancellation_flag_ =
-      std::make_unique<CancellationFlag>(true);
+  std::shared_ptr<CancellationFlag> default_cancellation_flag_ =
+      std::make_shared<CancellationFlag>(true);
 
   // An app lifecycle monitor for monitoring the app lifecycle state.
   std::unique_ptr<api::AppLifecycleMonitor> app_lifecycle_monitor_;
diff --git a/connections/implementation/client_proxy_test.cc b/connections/implementation/client_proxy_test.cc
index 5ea52f45..e0e1f757 100644
--- a/connections/implementation/client_proxy_test.cc
+++ b/connections/implementation/client_proxy_test.cc
@@ -14,6 +14,7 @@
 
 #include "connections/implementation/client_proxy.h"
 
+#include <atomic>
 #include <cstdint>
 #include <memory>
 #include <optional>
@@ -24,6 +25,7 @@
 #include "gmock/gmock.h"
 #include "protobuf-matchers/protocol-buffer-matchers.h"
 #include "gtest/gtest.h"
+#include "absl/base/thread_annotations.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
 #include "absl/time/clock.h"
@@ -38,6 +40,7 @@
 #include "connections/payload.h"
 #include "connections/status.h"
 #include "connections/strategy.h"
+#include "connections/v3/bandwidth_info.h"
 #include "connections/v3/connection_listening_options.h"
 #include "connections/v3/connection_result.h"
 #include "connections/v3/connections_device_provider.h"
@@ -53,6 +56,7 @@
 #include "internal/platform/medium_environment.h"
 #include "internal/platform/mutex.h"
 #include "internal/platform/mutex_lock.h"
+#include "internal/platform/single_thread_executor.h"
 #include "proto/connections_enums.pb.h"
 
 namespace nearby {
@@ -114,7 +118,7 @@ class FakeEventLogger : public ::nearby::analytics::MockEventLogger {
   }
 
   Mutex mutex_;
-  std::vector<ConnectionsLog> logs_;
+  std::vector<ConnectionsLog> logs_ ABSL_GUARDED_BY(mutex_);
 };
 
 class MockDeviceProvider : public nearby::NearbyDeviceProvider {
@@ -434,7 +438,7 @@ TEST_P(ClientProxyTest, CanCancelEndpoint) {
   // `CancellationFlag` pointers are passed to other classes in Nearby
   // Connections, and by using the pointers directly, we test their
   // consumption of `CancellationFlag` pointers.
-  CancellationFlag* cancellation_flag =
+  std::shared_ptr<CancellationFlag> cancellation_flag =
       client2()->GetCancellationFlag(advertising_endpoint.id);
 
   EXPECT_FALSE(
@@ -470,7 +474,7 @@ TEST_P(ClientProxyTest, CanCancelAllEndpoints) {
   // `CancellationFlag` pointers are passed to other classes in Nearby
   // Connections, and by using the pointers directly, we test their
   // consumption of `CancellationFlag` pointers.
-  CancellationFlag* cancellation_flag =
+  std::shared_ptr<CancellationFlag> cancellation_flag =
       client2()->GetCancellationFlag(advertising_endpoint.id);
 
   EXPECT_FALSE(
@@ -537,6 +541,26 @@ TEST_P(ClientProxyTest, CanCancelAllEndpointsWithDifferentEndpoint) {
   }
 }
 
+TEST_P(ClientProxyTest, GetCancellationFlagRace) {
+  std::string endpoint_id = "test_endpoint";
+  client1()->AddCancellationFlag(endpoint_id);
+
+  std::atomic<bool> run{true};
+  SingleThreadExecutor executor;
+  executor.Execute([&]() {
+    while (run) {
+      client1()->GetCancellationFlag(endpoint_id);
+    }
+  });
+
+  for (int i = 0; i < 10000; ++i) {
+    client1()->Reset();
+    client1()->AddCancellationFlag(endpoint_id);
+  }
+
+  run = false;
+}
+
 INSTANTIATE_TEST_SUITE_P(ParametrisedClientProxyTest, ClientProxyTest,
                          ::testing::ValuesIn(kTestCases));
 
diff --git a/connections/implementation/mediums/awdl_bwu_handler.cc b/connections/implementation/mediums/awdl_bwu_handler.cc
index edff8dda..8f8e5287 100644
--- a/connections/implementation/mediums/awdl_bwu_handler.cc
+++ b/connections/implementation/mediums/awdl_bwu_handler.cc
@@ -35,6 +35,7 @@
 #include "internal/base/masker.h"
 #include "internal/platform/awdl.h"
 #include "internal/platform/byte_array.h"
+#include "internal/platform/cancellation_flag.h"
 #include "internal/platform/count_down_latch.h"
 #include "internal/platform/expected.h"
 #include "internal/platform/implementation/psk_info.h"
@@ -149,9 +150,10 @@ AwdlBwuHandler::CreateUpgradedEndpointChannel(
             << service_name << ", service_type:" << service_type
             << ") for endpoint " << endpoint_id;
 
-  ErrorOr<AwdlSocket> socket_result =
-      awdl_medium_.Connect(upgrade_service_id, nsd_service_info, psk_info,
-                           client->GetCancellationFlag(endpoint_id));
+  std::shared_ptr<CancellationFlag> cancellation_flag =
+      client->GetCancellationFlag(endpoint_id);
+  ErrorOr<AwdlSocket> socket_result = awdl_medium_.Connect(
+      upgrade_service_id, nsd_service_info, psk_info, cancellation_flag.get());
   if (socket_result.has_error()) {
     LOG(ERROR) << "Failed to connect to the AWDL service (service_name:"
                << service_name << ", service_type:" << service_type
diff --git a/connections/implementation/mediums/ble/BUILD b/connections/implementation/mediums/ble/BUILD
index 619672c2..a4f1a0a4 100644
--- a/connections/implementation/mediums/ble/BUILD
+++ b/connections/implementation/mediums/ble/BUILD
@@ -70,7 +70,6 @@ cc_library(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
-        "@com_google_absl//absl/synchronization",
     ],
 )
 
diff --git a/connections/implementation/mediums/ble/ble_socket.cc b/connections/implementation/mediums/ble/ble_socket.cc
index 97c29e32..dc068c9c 100644
--- a/connections/implementation/mediums/ble/ble_socket.cc
+++ b/connections/implementation/mediums/ble/ble_socket.cc
@@ -207,22 +207,28 @@ Medium BleSocket::GetMediumLocked() const {
 }
 
 ExceptionOr<ByteArray> BleSocket::DispatchPacket() {
-  MutexLock lock(&mutex_);
-  if (!ble_input_stream_) {
-    return Exception::kFailed;
+  std::shared_ptr<BleInputStream> input_stream;
+  {
+    MutexLock lock(&mutex_);
+    if (!ble_input_stream_) {
+      return Exception::kFailed;
+    }
+    input_stream = ble_input_stream_;
   }
 
   ExceptionOr<ByteArray> read_bytes =
-      ble_input_stream_->Read(BlePacket::kServiceIdHashLength);
+      input_stream->Read(BlePacket::kServiceIdHashLength);
   while (read_bytes.ok()) {
     ByteArray read_bytes_result = read_bytes.result();
     if (BlePacket::IsControlPacketBytes(read_bytes_result)) {
-      ExceptionOr<ByteArray> handle_result = ProcessBleControlPacketLocked();
+      ExceptionOr<ByteArray> handle_result =
+          ProcessBleControlPacket(input_stream);
       if (!handle_result.ok()) {
         return handle_result;
       }
-      read_bytes = ble_input_stream_->Read(BlePacket::kServiceIdHashLength);
+      read_bytes = input_stream->Read(BlePacket::kServiceIdHashLength);
     } else {
+      MutexLock lock(&mutex_);
       if (read_bytes_result != service_id_hash_) {
         LOG(WARNING)
             << "Received data packet with incorrect service ID hash. Expected: "
@@ -239,20 +245,22 @@ ExceptionOr<ByteArray> BleSocket::DispatchPacket() {
 
 ExceptionOr<std::int32_t> BleSocket::ReadPayloadLength() {
   int payload_length = 0;
+  std::shared_ptr<BleInputStream> input_stream;
   {
     MutexLock lock(&mutex_);
     if (!ble_input_stream_) {
       return {Exception::kIo};
     }
-
-    ExceptionOr<ByteArray> read_bytes =
-        ble_input_stream_->Read(sizeof(std::int32_t));
-    if (!read_bytes.ok()) {
-      return read_bytes.exception();
-    }
-
-    payload_length = byte_utils::BytesToInt(std::move(read_bytes.result()));
+    input_stream = ble_input_stream_;
   }
+
+  ExceptionOr<ByteArray> read_bytes = input_stream->Read(sizeof(std::int32_t));
+  if (!read_bytes.ok()) {
+    return read_bytes.exception();
+  }
+
+  payload_length = byte_utils::BytesToInt(std::move(read_bytes.result()));
+
   Exception send_ack_result = SendPacketAcknowledgement(payload_length);
   if (!send_ack_result.Ok()) {
     LOG(WARNING) << "Failed to send packet acknowledgement.";
@@ -268,9 +276,10 @@ Exception BleSocket::WritePayloadLength(int payload_length) {
   return ble_output_stream_->WritePayloadLength(payload_length);
 }
 
-ExceptionOr<ByteArray> BleSocket::ProcessBleControlPacketLocked() {
+ExceptionOr<ByteArray> BleSocket::ProcessBleControlPacket(
+    std::shared_ptr<BleInputStream> input_stream) {
   // Read the first 4 bytes (packet block 1).
-  ExceptionOr<ByteArray> read_bytes = ble_input_stream_->Read(4);
+  ExceptionOr<ByteArray> read_bytes = input_stream->Read(4);
   if (!read_bytes.ok()) {
     return read_bytes;
   }
@@ -282,7 +291,7 @@ ExceptionOr<ByteArray> BleSocket::ProcessBleControlPacketLocked() {
   // Read the length from the 3rd byte of the packet block (0-indexed).
   int packet_block_2_size = packet_block_1.data()[3];
   // Read the left bytes for the packet block 2).
-  read_bytes = ble_input_stream_->Read(packet_block_2_size);
+  read_bytes = input_stream->Read(packet_block_2_size);
   if (!read_bytes.ok()) {
     return read_bytes;
   }
diff --git a/connections/implementation/mediums/ble/ble_socket.h b/connections/implementation/mediums/ble/ble_socket.h
index 4aba0ca1..85aa9fb0 100644
--- a/connections/implementation/mediums/ble/ble_socket.h
+++ b/connections/implementation/mediums/ble/ble_socket.h
@@ -375,8 +375,8 @@ class BleSocket final {
    * payload, the `ByteArray` may be empty. Returns an `Exception` if a
    * protocol error occurs or the read operation fails.
    */
-  ExceptionOr<ByteArray> ProcessBleControlPacketLocked()
-      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+  ExceptionOr<ByteArray> ProcessBleControlPacket(
+      std::shared_ptr<BleInputStream> input_stream);
 
   /**
    * Sends a raw L2CAP packet over the socket.
@@ -407,9 +407,9 @@ class BleSocket final {
   SingleThreadExecutor serial_executor_;
 
   const ByteArray service_id_hash_;
-  std::unique_ptr<mediums::BleInputStream> ble_input_stream_
+  std::shared_ptr<mediums::BleInputStream> ble_input_stream_
       ABSL_GUARDED_BY(mutex_) = nullptr;
-  std::unique_ptr<mediums::BleOutputStream> ble_output_stream_
+  std::shared_ptr<mediums::BleOutputStream> ble_output_stream_
       ABSL_GUARDED_BY(mutex_) = nullptr;
   nearby::BleSocket ble_socket_ ABSL_GUARDED_BY(mutex_) = nearby::BleSocket();
   nearby::BleL2capSocket l2cap_socket_ ABSL_GUARDED_BY(mutex_) =
diff --git a/connections/implementation/mediums/bluetooth_bwu_handler.cc b/connections/implementation/mediums/bluetooth_bwu_handler.cc
index 2a7e9075..0edaae99 100644
--- a/connections/implementation/mediums/bluetooth_bwu_handler.cc
+++ b/connections/implementation/mediums/bluetooth_bwu_handler.cc
@@ -18,17 +18,18 @@
 #include <string>
 #include <utility>
 
+#include "absl/base/nullability.h"
 #include "absl/functional/bind_front.h"
 #include "connections/implementation/base_bwu_handler.h"
 #include "connections/implementation/client_proxy.h"
 #include "connections/implementation/endpoint_channel.h"
 #include "connections/implementation/mediums/bluetooth_classic.h"
 #include "connections/implementation/mediums/bluetooth_endpoint_channel.h"
-#include "absl/base/nullability.h"
 #include "connections/implementation/mediums/bluetooth_radio.h"
 #include "connections/implementation/offline_frames.h"
 #include "internal/platform/bluetooth_adapter.h"
 #include "internal/platform/bluetooth_classic.h"
+#include "internal/platform/cancellation_flag.h"
 #include "internal/platform/expected.h"
 #include "internal/platform/logging.h"
 #include "internal/platform/mac_address.h"
@@ -89,8 +90,10 @@ BluetoothBwuHandler::CreateUpgradedEndpointChannel(
         OperationResultCode::CONNECTIVITY_BLUETOOTH_DEVICE_OBTAIN_FAILURE)};
   }
 
-  ErrorOr<BluetoothSocket> socket_result = bluetooth_medium_.Connect(
-      device, service_id, client->GetCancellationFlag(endpoint_id));
+  std::shared_ptr<CancellationFlag> cancellation_flag =
+      client->GetCancellationFlag(endpoint_id);
+  ErrorOr<BluetoothSocket> socket_result =
+      bluetooth_medium_.Connect(device, service_name, cancellation_flag.get());
   if (socket_result.has_error()) {
     LOG(ERROR)
         << "BluetoothBwuHandler failed to connect to the Bluetooth device ("
diff --git a/connections/implementation/mediums/bluetooth_bwu_handler_test.cc b/connections/implementation/mediums/bluetooth_bwu_handler_test.cc
index 4a5bac2c..55c6c021 100644
--- a/connections/implementation/mediums/bluetooth_bwu_handler_test.cc
+++ b/connections/implementation/mediums/bluetooth_bwu_handler_test.cc
@@ -44,10 +44,19 @@ constexpr absl::Duration kWaitDuration = absl::Milliseconds(1000);
 
 class BluetoothBwuTest : public testing::Test {
  protected:
-  BluetoothBwuTest() { env_.Start(); }
-  ~BluetoothBwuTest() override { env_.Stop(); }
+  BluetoothBwuTest() {
+    original_flags_ = FeatureFlags::GetInstance().GetFlags();
+    env_.Start();
+  }
+  ~BluetoothBwuTest() override {
+    FeatureFlags::GetMutableInstanceForTesting().SetFlags(original_flags_);
+    env_.Stop();
+  }
+
+  void RunSTACreateEndpointChannelTest(bool enable_cancellation);
 
   MediumEnvironment& env_{MediumEnvironment::Instance()};
+  FeatureFlags::Flags original_flags_;
 };
 
 TEST_F(BluetoothBwuTest, CanCreateBwuHandler) {
@@ -64,7 +73,12 @@ TEST_F(BluetoothBwuTest, CanCreateBwuHandler) {
   handler.reset();
 }
 
-TEST_F(BluetoothBwuTest, SoftAPBWUInit_STACreateEndpointChannel) {
+void BluetoothBwuTest::RunSTACreateEndpointChannelTest(
+    bool enable_cancellation) {
+  FeatureFlags::Flags flags = original_flags_;
+  flags.enable_cancellation_flag = enable_cancellation;
+  FeatureFlags::GetMutableInstanceForTesting().SetFlags(flags);
+
   CountDownLatch start_latch(1);
   CountDownLatch accept_latch(1);
   CountDownLatch end_latch(1);
@@ -73,6 +87,9 @@ TEST_F(BluetoothBwuTest, SoftAPBWUInit_STACreateEndpointChannel) {
   Mediums mediums_1, mediums_2;
   ExceptionOr<OfflineFrame> upgrade_frame;
 
+  EXPECT_TRUE(mediums_1.GetBluetoothRadio().Enable());
+  EXPECT_TRUE(mediums_2.GetBluetoothRadio().Enable());
+
   auto handler_1 = std::make_unique<BluetoothBwuHandler>(
       &mediums_1.GetBluetoothRadio(), &mediums_1.GetBluetoothClassic(),
       [&](ClientProxy* client,
@@ -113,7 +130,7 @@ TEST_F(BluetoothBwuTest, SoftAPBWUInit_STACreateEndpointChannel) {
         handler_2->CreateUpgradedEndpointChannel(&client_2, /*service_id=*/"A",
                                                  /*endpoint_id=*/"1",
                                                  bwu_frame.upgrade_path_info());
-    if (!FeatureFlags::GetInstance().GetFlags().enable_cancellation_flag) {
+    if (!enable_cancellation) {
       ASSERT_TRUE(result.has_value());
       std::unique_ptr<EndpointChannel> new_channel = std::move(result.value());
       EXPECT_TRUE(accept_latch.Await(kWaitDuration).result());
@@ -122,9 +139,9 @@ TEST_F(BluetoothBwuTest, SoftAPBWUInit_STACreateEndpointChannel) {
     } else {
       EXPECT_FALSE(result.has_value());
       EXPECT_TRUE(result.has_error());
-      EXPECT_EQ(
-          result.error().operation_result_code(),
-          OperationResultCode::CONNECTIVITY_BLUETOOTH_DEVICE_OBTAIN_FAILURE);
+      EXPECT_EQ(result.error().operation_result_code(),
+                OperationResultCode::
+                    CLIENT_CANCELLATION_CANCEL_BT_OUTGOING_CONNECTION);
       accept_latch.CountDown();
     }
     EXPECT_TRUE(mediums_2.GetBluetoothClassic().GetAddress().IsSet());
@@ -136,5 +153,15 @@ TEST_F(BluetoothBwuTest, SoftAPBWUInit_STACreateEndpointChannel) {
   EXPECT_TRUE(end_latch.Await(kWaitDuration).result());
 }
 
+TEST_F(BluetoothBwuTest,
+       SoftAPBWUInit_STACreateEndpointChannel_WithCancellation) {
+  RunSTACreateEndpointChannelTest(true);
+}
+
+TEST_F(BluetoothBwuTest,
+       SoftAPBWUInit_STACreateEndpointChannel_NoCancellation) {
+  RunSTACreateEndpointChannelTest(false);
+}
+
 }  // namespace connections
 }  // namespace nearby
diff --git a/connections/implementation/mediums/webrtc/BUILD b/connections/implementation/mediums/webrtc/BUILD
index c4305093..5b4f0b5b 100644
--- a/connections/implementation/mediums/webrtc/BUILD
+++ b/connections/implementation/mediums/webrtc/BUILD
@@ -208,6 +208,7 @@ cc_test(
     srcs = [
         "connection_flow_test.cc",
         "signaling_frames_test.cc",
+        "webrtc_bwu_handler_test.cc",
         "webrtc_impl_test.cc",
         "webrtc_socket_impl_test.cc",
     ],
@@ -223,11 +224,16 @@ cc_test(
         ":webrtc_impl",
         ":webrtc_medium",
         ":webrtc_socket_impl",
+        "//connections/implementation:bwu_handler",
+        "//connections/implementation:client_proxy",
+        "//connections/implementation:endpoint_channel",
+        "//connections/implementation:offline_frames",
         "//connections/implementation/mediums:webrtc",
         "//connections/implementation/mediums:webrtc_peer_id",
         "//connections/implementation/mediums:webrtc_socket",
         "//internal/platform:base",
         "//internal/platform:cancellation_flag",
+        "//internal/platform:logging",
         "//internal/platform:test_util",
         "//internal/platform:types",
         "//internal/platform/implementation:platform_impl",
diff --git a/connections/implementation/mediums/webrtc/webrtc_bwu_handler.cc b/connections/implementation/mediums/webrtc/webrtc_bwu_handler.cc
index 6c6c1dc0..1b2dac5d 100644
--- a/connections/implementation/mediums/webrtc/webrtc_bwu_handler.cc
+++ b/connections/implementation/mediums/webrtc/webrtc_bwu_handler.cc
@@ -29,6 +29,7 @@
 #include "connections/implementation/mediums/webrtc_socket.h"
 #include "connections/implementation/offline_frames.h"
 #include "connections/implementation/proto/offline_wire_formats.pb.h"
+#include "internal/platform/cancellation_flag.h"
 #include "internal/platform/expected.h"
 #include "internal/platform/implementation/webrtc_platform.h"
 #include "internal/platform/logging.h"
@@ -94,10 +95,11 @@ WebrtcBwuHandler::CreateUpgradedEndpointChannel(
             << peer_id.GetId() << ", location hint "
             << location_hint.location();
 
+  std::shared_ptr<CancellationFlag> cancellation_flag =
+      client->GetCancellationFlag(endpoint_id);
   ErrorOr<std::shared_ptr<mediums::WebRtcSocket>> socket_result =
       webrtc_.Connect(service_id, peer_id, location_hint,
-                      client->GetCancellationFlag(endpoint_id),
-                      client->GetWebRtcNonCellular());
+                      cancellation_flag.get(), client->GetWebRtcNonCellular());
   if (socket_result.has_error()) {
     LOG(ERROR) << "WebRtcBwuHandler failed to connect to remote peer ("
                << peer_id.GetId() << ") on endpoint " << endpoint_id
diff --git a/connections/implementation/mediums/webrtc/webrtc_bwu_handler_test.cc b/connections/implementation/mediums/webrtc/webrtc_bwu_handler_test.cc
new file mode 100644
index 00000000..fc407a65
--- /dev/null
+++ b/connections/implementation/mediums/webrtc/webrtc_bwu_handler_test.cc
@@ -0,0 +1,139 @@
+
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "connections/implementation/mediums/webrtc/webrtc_bwu_handler.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "gtest/gtest.h"
+#include "absl/time/time.h"
+#include "connections/implementation/bwu_handler.h"
+#include "connections/implementation/client_proxy.h"
+#include "connections/implementation/endpoint_channel.h"
+#include "connections/implementation/mediums/webrtc/webrtc_impl.h"
+#include "connections/implementation/offline_frames.h"
+#include "internal/platform/count_down_latch.h"
+#include "internal/platform/exception.h"
+#include "internal/platform/expected.h"
+#include "internal/platform/feature_flags.h"
+#include "internal/platform/logging.h"
+#include "internal/platform/medium_environment.h"
+#include "internal/platform/single_thread_executor.h"
+namespace nearby {
+namespace connections {
+namespace {
+using ::location::nearby::connections::OfflineFrame;
+using ::location::nearby::proto::connections::OperationResultCode;
+constexpr absl::Duration kWaitDuration = absl::Milliseconds(5000);
+class WebrtcBwuTest : public ::testing::Test {
+ protected:
+  WebrtcBwuTest() {
+    original_flags_ = FeatureFlags::GetInstance().GetFlags();
+    env_.Start({.webrtc_enabled = true});
+  }
+  ~WebrtcBwuTest() override {
+    FeatureFlags::GetMutableInstanceForTesting().SetFlags(original_flags_);
+    env_.Stop();
+  }
+  void RunCreateEndpointChannelTest(bool enable_cancellation);
+  MediumEnvironment& env_{MediumEnvironment::Instance()};
+  FeatureFlags::Flags original_flags_;
+};
+void WebrtcBwuTest::RunCreateEndpointChannelTest(bool enable_cancellation) {
+  FeatureFlags::Flags flags = original_flags_;
+  flags.enable_cancellation_flag = enable_cancellation;
+  FeatureFlags::GetMutableInstanceForTesting().SetFlags(flags);
+  CountDownLatch start_latch(1);
+  CountDownLatch accept_latch(1);
+  CountDownLatch end_latch(1);
+  ClientProxy client_1, client_2;
+  auto webrtc_1 = std::make_unique<mediums::WebRtcImpl>();
+  auto webrtc_2 = std::make_unique<mediums::WebRtcImpl>();
+  ExceptionOr<OfflineFrame> upgrade_frame;
+  std::unique_ptr<BwuHandler> handler_1 = std::make_unique<WebrtcBwuHandler>(
+      webrtc_1.get(),
+      [&](ClientProxy* client,
+          std::unique_ptr<BwuHandler::IncomingSocketConnection> connection) {
+        LOG(INFO) << "Handler 1 callback triggered";
+        accept_latch.CountDown();
+      });
+  std::unique_ptr<BwuHandler> handler_2 = std::make_unique<WebrtcBwuHandler>(
+      webrtc_2.get(),
+      [&](ClientProxy* client,
+          std::unique_ptr<BwuHandler::IncomingSocketConnection> connection) {
+        LOG(INFO) << "Handler 2 callback triggered";
+      });
+  // Server starts advertising.
+  SingleThreadExecutor server_executor;
+  server_executor.Execute([&]() {
+    std::string upgrade_frame_bytes =
+        handler_1->InitializeUpgradedMediumForEndpoint(
+            &client_1, /*upgrade_service_id=*/"A", /*endpoint_id=*/"1");
+    EXPECT_FALSE(upgrade_frame_bytes.empty());
+    upgrade_frame = parser::FromBytes(upgrade_frame_bytes);
+    start_latch.CountDown();
+  });
+  // Client connects.
+  EXPECT_TRUE(start_latch.Await(kWaitDuration).result());
+  if (enable_cancellation) {
+    client_2.AddCancellationFlag(/*endpoint_id=*/"1");
+    client_2.GetCancellationFlag(/*endpoint_id=*/"1")->Cancel();
+  }
+  SingleThreadExecutor client_executor;
+  client_executor.Execute([&]() {
+    auto bwu_frame =
+        upgrade_frame.result().v1().bandwidth_upgrade_negotiation();
+    auto result = handler_2->CreateUpgradedEndpointChannel(
+        &client_2, /*service_id=*/"A",
+        /*endpoint_id=*/"1", bwu_frame.upgrade_path_info());
+    if (!enable_cancellation) {
+      ASSERT_TRUE(result.has_value());
+      std::unique_ptr<EndpointChannel> new_channel = std::move(result.value());
+      EXPECT_TRUE(accept_latch.Await(kWaitDuration).result());
+      EXPECT_EQ(new_channel->GetMedium(),
+                location::nearby::proto::connections::Medium::WEB_RTC);
+    } else {
+      EXPECT_FALSE(result.has_value());
+      EXPECT_TRUE(result.has_error());
+      EXPECT_EQ(result.error().operation_result_code(),
+                OperationResultCode::
+                    CLIENT_CANCELLATION_CANCEL_WEB_RTC_OUTGOING_CONNECTION);
+      accept_latch.CountDown();
+    }
+    handler_1->RevertResponderState(/*service_id=*/"A");
+    end_latch.CountDown();
+  });
+  EXPECT_TRUE(accept_latch.Await(kWaitDuration).result());
+  EXPECT_TRUE(end_latch.Await(kWaitDuration).result());
+}
+TEST_F(WebrtcBwuTest, CanCreateBwuHandler) {
+  auto webrtc = std::make_unique<mediums::WebRtcImpl>();
+  std::unique_ptr<BwuHandler> handler = std::make_unique<WebrtcBwuHandler>(
+      webrtc.get(),
+      [](ClientProxy* client,
+         std::unique_ptr<BwuHandler::IncomingSocketConnection> connection) {});
+  EXPECT_EQ(handler->GetUpgradeMedium(),
+            location::nearby::proto::connections::Medium::WEB_RTC);
+}
+TEST_F(WebrtcBwuTest, CreateEndpointChannel_WithCancellation) {
+  RunCreateEndpointChannelTest(true);
+}
+TEST_F(WebrtcBwuTest, CreateEndpointChannel_NoCancellation) {
+  RunCreateEndpointChannelTest(false);
+}
+}  // namespace
+}  // namespace connections
+}  // namespace nearby
diff --git a/connections/implementation/mediums/wifi_direct_bwu_handler.cc b/connections/implementation/mediums/wifi_direct_bwu_handler.cc
index 5072e922..b9925b84 100644
--- a/connections/implementation/mediums/wifi_direct_bwu_handler.cc
+++ b/connections/implementation/mediums/wifi_direct_bwu_handler.cc
@@ -19,16 +19,17 @@
 #include <string>
 #include <utility>
 
+#include "absl/base/nullability.h"
 #include "absl/functional/bind_front.h"
 #include "connections/implementation/base_bwu_handler.h"
 #include "connections/implementation/client_proxy.h"
 #include "connections/implementation/endpoint_channel.h"
-#include "absl/base/nullability.h"
 #include "connections/implementation/mediums/wifi_direct.h"
 #include "connections/implementation/mediums/wifi_direct_endpoint_channel.h"
 #include "connections/implementation/offline_frames.h"
 #include "connections/strategy.h"
 #include "internal/base/masker.h"
+#include "internal/platform/cancellation_flag.h"
 #include "internal/platform/expected.h"
 #include "internal/platform/logging.h"
 #include "internal/platform/wifi_credential.h"
@@ -173,8 +174,10 @@ WifiDirectBwuHandler::CreateUpgradedEndpointChannel(
         OperationResultCode::CONNECTIVITY_WIFI_DIRECT_INVALID_CREDENTIAL)};
   }
 
+  std::shared_ptr<CancellationFlag> cancellation_flag =
+      client->GetCancellationFlag(endpoint_id);
   ErrorOr<WifiDirectSocket> socket_result = wifi_direct_medium_.Connect(
-      service_id, gateway, port, client->GetCancellationFlag(endpoint_id));
+      service_id, gateway, port, cancellation_flag.get());
   if (socket_result.has_error()) {
     LOG(ERROR)
         << "WifiDirectBwuHandler failed to connect to the WifiDirect service("
diff --git a/connections/implementation/mediums/wifi_hotspot_bwu_handler.cc b/connections/implementation/mediums/wifi_hotspot_bwu_handler.cc
index 77d14c0d..d31b8a15 100644
--- a/connections/implementation/mediums/wifi_hotspot_bwu_handler.cc
+++ b/connections/implementation/mediums/wifi_hotspot_bwu_handler.cc
@@ -38,6 +38,7 @@
 #include "connections/implementation/proto/offline_wire_formats.pb.h"
 #include "connections/strategy.h"
 #include "internal/base/masker.h"
+#include "internal/platform/cancellation_flag.h"
 #include "internal/platform/expected.h"
 #include "internal/platform/implementation/wifi_utils.h"
 #include "internal/platform/logging.h"
@@ -220,9 +221,11 @@ WifiHotspotBwuHandler::CreateUpgradedEndpointChannel(
                       CONNECTIVITY_WIFI_HOTSPOT_LEGACY_STA_CONNECTION_FAILURE)};
   }
 
+  std::shared_ptr<CancellationFlag> cancellation_flag =
+      client->GetCancellationFlag(endpoint_id);
   ErrorOr<WifiHotspotSocket> socket_result = wifi_hotspot_medium_.Connect(
       service_id, hotspot_credentials.GetAddressCandidates(),
-      client->GetCancellationFlag(endpoint_id));
+      cancellation_flag.get());
   if (socket_result.has_error()) {
     LOG(ERROR) << "WifiHotspotBwuHandler failed to connect to the WifiHotspot "
                   "service for endpoint "
diff --git a/connections/implementation/mediums/wifi_lan_bwu_handler.cc b/connections/implementation/mediums/wifi_lan_bwu_handler.cc
index e29e2f99..fca9ecef 100644
--- a/connections/implementation/mediums/wifi_lan_bwu_handler.cc
+++ b/connections/implementation/mediums/wifi_lan_bwu_handler.cc
@@ -20,14 +20,15 @@
 #include <utility>
 #include <vector>
 
+#include "absl/base/nullability.h"
 #include "absl/functional/bind_front.h"
 #include "connections/implementation/base_bwu_handler.h"
 #include "connections/implementation/client_proxy.h"
 #include "connections/implementation/endpoint_channel.h"
-#include "absl/base/nullability.h"
 #include "connections/implementation/mediums/wifi_lan.h"
 #include "connections/implementation/mediums/wifi_lan_endpoint_channel.h"
 #include "connections/implementation/offline_frames.h"
+#include "internal/platform/cancellation_flag.h"
 #include "internal/platform/expected.h"
 #include "internal/platform/implementation/upgrade_address_info.h"
 #include "internal/platform/logging.h"
@@ -93,9 +94,10 @@ WifiLanBwuHandler::CreateUpgradedEndpointChannel(
     VLOG(1) << "WifiLanBwuHandler is attempting to connect to available "
                "WifiLan service (" << address_candidate << ") for endpoint "
             << endpoint_id;
-    ErrorOr<WifiLanSocket> socket_result =
-        wifi_lan_medium_.Connect(service_id, address_candidate,
-                                 client->GetCancellationFlag(endpoint_id));
+    std::shared_ptr<CancellationFlag> cancellation_flag =
+        client->GetCancellationFlag(endpoint_id);
+    ErrorOr<WifiLanSocket> socket_result = wifi_lan_medium_.Connect(
+        service_id, address_candidate, cancellation_flag.get());
     if (socket_result.has_error()) {
       LOG(ERROR)
           << "WifiLanBwuHandler failed to connect to the WifiLan service ("
diff --git a/connections/implementation/p2p_cluster_pcp_handler.cc b/connections/implementation/p2p_cluster_pcp_handler.cc
index 76bbe1bc..9f530d60 100644
--- a/connections/implementation/p2p_cluster_pcp_handler.cc
+++ b/connections/implementation/p2p_cluster_pcp_handler.cc
@@ -67,6 +67,7 @@
 #include "internal/platform/bluetooth_adapter.h"
 #include "internal/platform/bluetooth_classic.h"
 #include "internal/platform/byte_array.h"
+#include "internal/platform/cancellation_flag.h"
 #include "internal/platform/expected.h"
 #include "internal/platform/implementation/platform.h"
 #include "internal/platform/logging.h"
@@ -2015,9 +2016,10 @@ BasePcpHandler::ConnectImplResult P2pClusterPcpHandler::BluetoothConnectImpl(
           << endpoint->endpoint_id << ") over Bluetooth Classic.";
   BluetoothDevice& device = endpoint->bluetooth_device;
 
+  std::shared_ptr<CancellationFlag> cancellation_flag =
+      client->GetCancellationFlag(endpoint->endpoint_id);
   ErrorOr<BluetoothSocket> bluetooth_socket_result = bluetooth_medium_.Connect(
-      device, endpoint->service_id,
-      client->GetCancellationFlag(endpoint->endpoint_id));
+      device, endpoint->service_id, cancellation_flag.get());
   if (bluetooth_socket_result.has_error()) {
     LOG(ERROR)
         << "In BluetoothConnectImpl(), failed to connect to Bluetooth device "
@@ -2386,15 +2388,17 @@ BasePcpHandler::ConnectImplResult P2pClusterPcpHandler::BleConnectImpl(
           << " is attempting to connect to (" << peripheral.ToReadableString()
           << ") over BLE.";
 
+  std::shared_ptr<CancellationFlag> cancellation_flag =
+      client->GetCancellationFlag(endpoint->endpoint_id);
+
   if (NearbyFlags::GetInstance().GetBoolFlag(
           config_package_nearby::nearby_connections_feature::kEnableBleL2cap) &&
       peripheral.GetPsm() !=
           mediums::BleAdvertisementHeader::kDefaultPsmValue) {
     if (refactor_ble_l2cap) {
       ErrorOr<std::unique_ptr<mediums::BleSocket>> ble_l2cap_socket_result =
-          ble_medium_.ConnectOverL2cap2(
-              endpoint->service_id, peripheral,
-              client->GetCancellationFlag(endpoint->endpoint_id));
+          ble_medium_.ConnectOverL2cap2(endpoint->service_id, peripheral,
+                                        cancellation_flag.get());
       if (!ble_l2cap_socket_result.has_error()) {
         LOG(INFO) << "In BleV2ConnectImpl(), connected to Ble L2CAP device "
                   << absl::BytesToHexString(peripheral.GetId().data())
@@ -2416,9 +2420,8 @@ BasePcpHandler::ConnectImplResult P2pClusterPcpHandler::BleConnectImpl(
       }
     } else {
       ErrorOr<BleL2capSocket> ble_l2cap_socket_result =
-          ble_medium_.ConnectOverL2cap(
-              endpoint->service_id, peripheral,
-              client->GetCancellationFlag(endpoint->endpoint_id));
+          ble_medium_.ConnectOverL2cap(endpoint->service_id, peripheral,
+                                       cancellation_flag.get());
       if (!ble_l2cap_socket_result.has_error()) {
         LOG(INFO) << "In BleConnectImpl(), connected to Ble L2CAP device "
                   << absl::BytesToHexString(peripheral.GetId().data())
@@ -2444,9 +2447,8 @@ BasePcpHandler::ConnectImplResult P2pClusterPcpHandler::BleConnectImpl(
   std::unique_ptr<BleEndpointChannel> channel = nullptr;
   if (refactor_ble_l2cap) {
     ErrorOr<std::unique_ptr<mediums::BleSocket>> ble_socket_result =
-        ble_medium_.Connect2(
-            endpoint->service_id, peripheral,
-            client->GetCancellationFlag(endpoint->endpoint_id));
+        ble_medium_.Connect2(endpoint->service_id, peripheral,
+                             cancellation_flag.get());
     if (ble_socket_result.has_error()) {
       LOG(ERROR) << "In BleConnectImpl(), failed to connect to BLE device "
                  << absl::BytesToHexString(peripheral.GetId().data())
@@ -2461,9 +2463,8 @@ BasePcpHandler::ConnectImplResult P2pClusterPcpHandler::BleConnectImpl(
         endpoint->service_id, /*channel_name=*/endpoint->endpoint_id,
         std::move(ble_socket_result.value()));
   } else {
-    ErrorOr<BleSocket> ble_socket_result =
-        ble_medium_.Connect(endpoint->service_id, peripheral,
-                            client->GetCancellationFlag(endpoint->endpoint_id));
+    ErrorOr<BleSocket> ble_socket_result = ble_medium_.Connect(
+        endpoint->service_id, peripheral, cancellation_flag.get());
     if (ble_socket_result.has_error()) {
       LOG(ERROR) << "In BleConnectImpl(), failed to connect to BLE device "
                  << absl::BytesToHexString(peripheral.GetId().data())
@@ -2737,9 +2738,10 @@ BasePcpHandler::ConnectImplResult P2pClusterPcpHandler::AwdlConnectImpl(
   LOG(INFO) << "Client " << client->GetClientId()
             << " is attempting to connect to endpoint(id="
             << endpoint->endpoint_id << ") over Awdl.";
-  ErrorOr<AwdlSocket> socket_result =
-      awdl_medium_.Connect(endpoint->service_id, endpoint->service_info,
-                           client->GetCancellationFlag(endpoint->endpoint_id));
+  std::shared_ptr<CancellationFlag> cancellation_flag =
+      client->GetCancellationFlag(endpoint->endpoint_id);
+  ErrorOr<AwdlSocket> socket_result = awdl_medium_.Connect(
+      endpoint->service_id, endpoint->service_info, cancellation_flag.get());
   if (socket_result.has_error()) {
     LOG(ERROR) << "In AwdlConnectImpl(), failed to connect to service "
                << endpoint->service_info.GetServiceName()
@@ -2773,9 +2775,10 @@ BasePcpHandler::ConnectImplResult P2pClusterPcpHandler::WifiLanConnectImpl(
   LOG(INFO) << "Client " << client->GetClientId()
             << " is attempting to connect to endpoint(id="
             << endpoint->endpoint_id << ") over WifiLan.";
+  std::shared_ptr<CancellationFlag> cancellation_flag =
+      client->GetCancellationFlag(endpoint->endpoint_id);
   ErrorOr<WifiLanSocket> socket_result = wifi_lan_medium_.Connect(
-      endpoint->service_id, endpoint->service_info,
-      client->GetCancellationFlag(endpoint->endpoint_id));
+      endpoint->service_id, endpoint->service_info, cancellation_flag.get());
   if (socket_result.has_error()) {
     LOG(ERROR) << "In WifiLanConnectImpl(), failed to connect to service "
                << endpoint->service_info.GetServiceName()
diff --git a/connections/implementation/p2p_cluster_pcp_handler_test.cc b/connections/implementation/p2p_cluster_pcp_handler_test.cc
index 956856a7..5f6c813e 100644
--- a/connections/implementation/p2p_cluster_pcp_handler_test.cc
+++ b/connections/implementation/p2p_cluster_pcp_handler_test.cc
@@ -88,11 +88,15 @@ class P2pClusterPcpHandlerTest : public testing::Test {
     LOG(INFO) << "SetUp: begin";
     NearbyFlags::GetInstance().OverrideBoolFlagValue(
         config_package_nearby::nearby_connections_feature::kEnableAwdl, true);
-    SetBleExtendedAdvertisementsAvailable(true);
+    SetBleExtendedAdvertisementsAvailable(false);
+  }
+
+  void TearDown() override {
+    NearbyFlags::GetInstance().ResetOverridedValues();
   }
 
   void SetBleExtendedAdvertisementsAvailable(bool available) {
-    env_.SetBleExtendedAdvertisementsAvailable(false);
+    env_.SetBleExtendedAdvertisementsAvailable(available);
   }
 
   AdvertisingOptions GetBluetoothOnlyAdvertisingOptions() {
@@ -141,6 +145,8 @@ class P2pClusterPcpHandlerTest : public testing::Test {
     return ByteArray(reinterpret_cast<char*>(bytes), 6);
   }
 
+  void RunCanConnectHelper(BooleanMediumSelector selector);
+
   ClientProxy client_a_;
   ClientProxy client_b_;
   ClientProxy client_c_;
@@ -148,6 +154,130 @@ class P2pClusterPcpHandlerTest : public testing::Test {
   MediumEnvironment& env_{MediumEnvironment::Instance()};
 };
 
+void P2pClusterPcpHandlerTest::RunCanConnectHelper(
+    BooleanMediumSelector selector) {
+  env_.Start();
+  std::string endpoint_name_a{"endpoint_name"};
+  Mediums mediums_a;
+  Mediums mediums_b;
+  BluetoothRadio& radio_a = mediums_a.GetBluetoothRadio();
+  BluetoothRadio& radio_b = mediums_b.GetBluetoothRadio();
+  radio_a.GetBluetoothAdapter().SetName("BT Device A");
+  radio_b.GetBluetoothAdapter().SetName("BT Device B");
+  EndpointChannelManager ecm_a;
+  EndpointChannelManager ecm_b;
+  EndpointManager em_a(&ecm_a);
+  EndpointManager em_b(&ecm_b);
+  BwuManager bwu_a(mediums_a, em_a, ecm_a, {},
+                   {.allow_upgrade_to = {.bluetooth = true}});
+  BwuManager bwu_b(mediums_b, em_b, ecm_b, {},
+                   {.allow_upgrade_to = {.bluetooth = true}});
+  InjectedBluetoothDeviceStore ibds_a;
+  InjectedBluetoothDeviceStore ibds_b;
+  P2pClusterPcpHandler handler_a(&mediums_a, &em_a, &ecm_a, &bwu_a, ibds_a);
+  P2pClusterPcpHandler handler_b(&mediums_b, &em_b, &ecm_b, &bwu_b, ibds_b);
+  CountDownLatch discover_latch(1);
+  CountDownLatch connect_latch(2);
+  struct DiscoveredInfo {
+    std::string endpoint_id;
+    ByteArray endpoint_info;
+    std::string service_id;
+  } discovered;
+
+  // Build options locally using passed selector!
+  AdvertisingOptions advertising_options = {{Strategy::kP2pCluster, selector}};
+  DiscoveryOptions discovery_options = {{Strategy::kP2pCluster, selector}};
+  ConnectionOptions connection_options = {{Strategy::kP2pCluster, selector}};
+
+  EXPECT_EQ(
+      handler_a.StartAdvertising(
+          &client_a_, service_id_, advertising_options,
+          {
+              .endpoint_info = ByteArray{endpoint_name_a},
+              .listener =
+                  {
+                      .initiated_cb =
+                          [&connect_latch](const std::string& endpoint_id,
+                                           const ConnectionResponseInfo& info) {
+                            LOG(INFO)
+                                << "StartAdvertising: initiated_cb called";
+                            connect_latch.CountDown();
+                          },
+                  },
+          }),
+      Status{Status::kSuccess});
+  EXPECT_EQ(handler_b.StartDiscovery(
+                &client_b_, service_id_, discovery_options,
+                {
+                    .endpoint_found_cb =
+                        [&discover_latch, &discovered](
+                            const std::string& endpoint_id,
+                            const ByteArray& endpoint_info,
+                            const std::string& service_id) {
+                          LOG(INFO) << "Device discovered: id=" << endpoint_id
+                                    << ", endpoint_info="
+                                    << std::string{endpoint_info};
+                          discovered = {
+                              .endpoint_id = endpoint_id,
+                              .endpoint_info = endpoint_info,
+                              .service_id = service_id,
+                          };
+                          discover_latch.CountDown();
+                        },
+                }),
+            Status{Status::kSuccess});
+
+  EXPECT_TRUE(discover_latch.Await(absl::Milliseconds(1000)).result());
+  EXPECT_EQ(endpoint_name_a, std::string{discovered.endpoint_info});
+
+  const std::string kBssid = "34:36:3B:C7:8C:71";
+  const std::int32_t kFreq = 5200;
+
+  connection_options.connection_info.supports_5_ghz = true;
+  connection_options.connection_info.bssid = kBssid;
+  connection_options.connection_info.ap_frequency = kFreq;
+
+  client_b_.AddCancellationFlag(discovered.endpoint_id);
+  handler_b.RequestConnection(
+      &client_b_, discovered.endpoint_id,
+      {.endpoint_info = discovered.endpoint_info,
+       .listener =
+           {
+               .initiated_cb =
+                   [&connect_latch](const std::string& endpoint_id,
+                                    const ConnectionResponseInfo& info) {
+                     LOG(INFO) << "RequestConnection: initiated_cb called";
+                     connect_latch.CountDown();
+                   },
+           }},
+      connection_options);
+  std::string client_b_local_endpoint = client_b_.GetLocalEndpointId();
+
+  EXPECT_TRUE(connect_latch.Await(absl::Milliseconds(1000)).result());
+  EXPECT_TRUE(client_b_.Is5GHzSupported(discovered.endpoint_id));
+  EXPECT_EQ(client_b_.GetBssid(discovered.endpoint_id), kBssid);
+  EXPECT_EQ(client_b_.GetApFrequency(discovered.endpoint_id), kFreq);
+  // When connection is established, EndpointManager will setup KeepAliveManager
+  // loop. When it fails, the connection will be dismantled. Since this a unit
+  // test, KeepAliveManager won't be really up. The disconnection may happen
+  // before the following check, which cause the check fail. So we check the
+  // connection status first.
+  if (client_b_.IsConnectedToEndpoint(discovered.endpoint_id)) {
+    EXPECT_EQ(client_a_.Is5GHzSupported(client_b_local_endpoint),
+              mediums_b.GetWifi().GetCapability().supports_5_ghz);
+    EXPECT_EQ(client_a_.GetBssid(client_b_local_endpoint),
+              mediums_b.GetWifi().GetInformation().bssid);
+    EXPECT_EQ(client_a_.GetApFrequency(client_b_local_endpoint),
+              mediums_b.GetWifi().GetInformation().ap_frequency);
+  }
+
+  handler_a.StopAdvertising(&client_a_);
+  handler_b.StopDiscovery(&client_b_);
+  bwu_a.Shutdown();
+  bwu_b.Shutdown();
+  env_.Stop();
+}
+
 TEST_F(P2pClusterPcpHandlerTest, NoBluetoothDiscoveryWhenRadioIsOff) {
   env_.Start();
   Mediums mediums;
@@ -231,7 +361,8 @@ TEST_F(P2pClusterPcpHandlerTest,
 }
 
 class P2pClusterPcpHandlerTestWithParam
-    : public testing::TestWithParam</*mediums=*/BooleanMediumSelector> {
+    : public P2pClusterPcpHandlerTest,
+      public ::testing::WithParamInterface</*mediums=*/BooleanMediumSelector> {
  protected:
   void SetUp() override {
     LOG(INFO) << "SetUp: begin";
@@ -264,9 +395,6 @@ class P2pClusterPcpHandlerTestWithParam
     LOG(INFO) << "SetUp: end";
   }
 
-  ClientProxy client_a_;
-  ClientProxy client_b_;
-  std::string service_id_{"service"};
   ConnectionOptions connection_options_{
       {
           Strategy::kP2pCluster,
@@ -285,7 +413,6 @@ class P2pClusterPcpHandlerTestWithParam
           GetParam(),
       },
   };
-  MediumEnvironment& env_{MediumEnvironment::Instance()};
 };
 
 TEST_P(P2pClusterPcpHandlerTestWithParam, CanConstructOne) {
@@ -900,128 +1027,21 @@ TEST_P(P2pClusterPcpHandlerTestWithParam,
 }
 
 TEST_P(P2pClusterPcpHandlerTestWithParam, CanConnect) {
-  env_.Start();
-  std::string endpoint_name_a{"endpoint_name"};
-  Mediums mediums_a;
-  Mediums mediums_b;
-  BluetoothRadio& radio_a = mediums_a.GetBluetoothRadio();
-  BluetoothRadio& radio_b = mediums_b.GetBluetoothRadio();
-  radio_a.GetBluetoothAdapter().SetName("BT Device A");
-  radio_b.GetBluetoothAdapter().SetName("BT Device B");
-  EndpointChannelManager ecm_a;
-  EndpointChannelManager ecm_b;
-  EndpointManager em_a(&ecm_a);
-  EndpointManager em_b(&ecm_b);
-  BwuManager bwu_a(mediums_a, em_a, ecm_a, {},
-                   {.allow_upgrade_to = {.bluetooth = true}});
-  BwuManager bwu_b(mediums_b, em_b, ecm_b, {},
-                   {.allow_upgrade_to = {.bluetooth = true}});
-  InjectedBluetoothDeviceStore ibds_a;
-  InjectedBluetoothDeviceStore ibds_b;
-  P2pClusterPcpHandler handler_a(&mediums_a, &em_a, &ecm_a, &bwu_a, ibds_a);
-  P2pClusterPcpHandler handler_b(&mediums_b, &em_b, &ecm_b, &bwu_b, ibds_b);
-  CountDownLatch discover_latch(1);
-  CountDownLatch connect_latch(2);
-  struct DiscoveredInfo {
-    std::string endpoint_id;
-    ByteArray endpoint_info;
-    std::string service_id;
-  } discovered;
-  EXPECT_EQ(
-      handler_a.StartAdvertising(
-          &client_a_, service_id_, advertising_options_,
-          {
-              .endpoint_info = ByteArray{endpoint_name_a},
-              .listener =
-                  {
-                      .initiated_cb =
-                          [&connect_latch](const std::string& endpoint_id,
-                                           const ConnectionResponseInfo& info) {
-                            LOG(INFO)
-                                << "StartAdvertising: initiated_cb called";
-                            connect_latch.CountDown();
-                          },
-                  },
-          }),
-      Status{Status::kSuccess});
-  EXPECT_EQ(handler_b.StartDiscovery(
-                &client_b_, service_id_, discovery_options_,
-                {
-                    .endpoint_found_cb =
-                        [&discover_latch, &discovered](
-                            const std::string& endpoint_id,
-                            const ByteArray& endpoint_info,
-                            const std::string& service_id) {
-                          LOG(INFO) << "Device discovered: id=" << endpoint_id
-                                    << ", endpoint_info="
-                                    << std::string{endpoint_info};
-                          discovered = {
-                              .endpoint_id = endpoint_id,
-                              .endpoint_info = endpoint_info,
-                              .service_id = service_id,
-                          };
-                          discover_latch.CountDown();
-                        },
-                }),
-            Status{Status::kSuccess});
-
-  EXPECT_TRUE(discover_latch.Await(absl::Milliseconds(1000)).result());
-  EXPECT_EQ(endpoint_name_a, std::string{discovered.endpoint_info});
-
-  const std::string kBssid = "34:36:3B:C7:8C:71";
-  const std::int32_t kFreq = 5200;
-
-  connection_options_.connection_info.supports_5_ghz = true;
-  connection_options_.connection_info.bssid = kBssid;
-  connection_options_.connection_info.ap_frequency = kFreq;
-
-  client_b_.AddCancellationFlag(discovered.endpoint_id);
-  handler_b.RequestConnection(
-      &client_b_, discovered.endpoint_id,
-      {.endpoint_info = discovered.endpoint_info,
-       .listener =
-           {
-               .initiated_cb =
-                   [&connect_latch](const std::string& endpoint_id,
-                                    const ConnectionResponseInfo& info) {
-                     LOG(INFO) << "RequestConnection: initiated_cb called";
-                     connect_latch.CountDown();
-                   },
-           }},
-      connection_options_);
-  std::string client_b_local_endpoint = client_b_.GetLocalEndpointId();
-
-  EXPECT_TRUE(connect_latch.Await(absl::Milliseconds(1000)).result());
-  EXPECT_TRUE(client_b_.Is5GHzSupported(discovered.endpoint_id));
-  EXPECT_EQ(client_b_.GetBssid(discovered.endpoint_id), kBssid);
-  EXPECT_EQ(client_b_.GetApFrequency(discovered.endpoint_id), kFreq);
-  // When connection is established, EndpointManager will setup KeepAliveManager
-  // loop. When it fails, the connection will be dismantled. Since this a unit
-  // test, KeepAliveManager won't be really up. The disconnection may happen
-  // before the following check, which cause the check fail. So we check the
-  // connection status first.
-  if (client_b_.IsConnectedToEndpoint(discovered.endpoint_id)) {
-    EXPECT_EQ(client_a_.Is5GHzSupported(client_b_local_endpoint),
-              mediums_b.GetWifi().GetCapability().supports_5_ghz);
-    EXPECT_EQ(client_a_.GetBssid(client_b_local_endpoint),
-              mediums_b.GetWifi().GetInformation().bssid);
-    EXPECT_EQ(client_a_.GetApFrequency(client_b_local_endpoint),
-              mediums_b.GetWifi().GetInformation().ap_frequency);
-  }
-
-  handler_a.StopAdvertising(&client_a_);
-  handler_b.StopDiscovery(&client_b_);
-  bwu_a.Shutdown();
-  bwu_b.Shutdown();
-  env_.Stop();
+  RunCanConnectHelper(GetParam());
 }
 
 TEST_P(P2pClusterPcpHandlerTestWithParam, CanConnectWithDctEnabled) {
   env_.Start();
+  // DCT advertisement truncates the device name to 7 bytes.
+  // "Test device" (11 bytes) -> "Test de" (7 bytes).
+  // The endpoint info is constructed by advertisements::BuildEndpointInfo which
+  // adds some overhead.
+  // For DCT, it seems to be 18 bytes prefix + truncated device name.
+  // 18 + 7 = 25 bytes.
   ByteArray endpoint_info_a{
-      "\x22\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0b"
-      "\x54\x65\x73\x74\x20\x64\x65\x76\x69\x63\x65",
-      29};
+      "\x22\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f\x07"
+      "Test de",
+      25};
   ClientProxy client_a;
   ClientProxy client_b;
 
@@ -1783,5 +1803,47 @@ INSTANTIATE_TEST_SUITE_P(ParametrisedPcpHandlerTest,
                          P2pClusterPcpHandlerTestWithParam,
                          ::testing::ValuesIn(kTestCases));
 
+TEST_F(P2pClusterPcpHandlerTest, BleConnect_L2cap_Refactor) {
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::kEnableBleL2cap, true);
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::kRefactorBleL2cap,
+      true);
+
+  RunCanConnectHelper({.ble = true});
+}
+
+TEST_F(P2pClusterPcpHandlerTest, BleConnect_NoL2cap_Refactor) {
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::kEnableBleL2cap,
+      false);
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::kRefactorBleL2cap,
+      true);
+
+  RunCanConnectHelper({.ble = true});
+}
+
+TEST_F(P2pClusterPcpHandlerTest, BleConnect_NoL2cap_NoRefactor) {
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::kEnableBleL2cap,
+      false);
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::kRefactorBleL2cap,
+      false);
+
+  RunCanConnectHelper({.ble = true});
+}
+
+TEST_F(P2pClusterPcpHandlerTest, BleConnect_L2cap_NoRefactor) {
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::kEnableBleL2cap, true);
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::kRefactorBleL2cap,
+      false);
+
+  RunCanConnectHelper({.ble = true});
+}
+
 }  // namespace
 }  // namespace nearby::connections