Implement dynamic role switching for bandwidth upgrades.

PiperOrigin-RevId: 933898362

connections/implementation/BUILD

@@ -349,6 +349,7 @@ cc_test(
         "@com_github_protobuf_matchers//protobuf-matchers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/time",
         "@com_google_googletest//:gtest_main",
     ],
 )

connections/implementation/base_endpoint_channel.cc

@@ -418,6 +418,15 @@ int BaseEndpointChannel::GetFrequency() const { return frequency_; }
 // Returns the try count of this EndpointChannel.
 int BaseEndpointChannel::GetTryCount() const { return try_count_; }
 
+void BaseEndpointChannel::SetLocalEndpointId(
+    const std::string& local_endpoint_id) {
+  local_endpoint_id_ = local_endpoint_id;
+}
+
+std::string BaseEndpointChannel::GetLocalEndpointId() const {
+  return local_endpoint_id_;
+}
+
 int BaseEndpointChannel::GetMaxAllowedReadBytes() const {
   int64_t max_allowed_read_bytes = NearbyFlags::GetInstance().GetInt64Flag(
       config_package_nearby::nearby_connections_feature::

connections/implementation/base_endpoint_channel.h

@@ -81,6 +81,8 @@ class BaseEndpointChannel : public EndpointChannel {
   uint32_t GetNextKeepAliveSeqNo() const override;
   void SetAnalyticsRecorder(analytics::AnalyticsRecorder* analytics_recorder,
                             const std::string& endpoint_id) override;
+  void SetLocalEndpointId(const std::string& local_endpoint_id) override;
+  std::string GetLocalEndpointId() const override;
 
   // Reads a complete packet from the underlying medium.
   virtual ExceptionOr<ByteArray> DispatchPacket() {
@@ -166,6 +168,7 @@ class BaseEndpointChannel : public EndpointChannel {
 
   analytics::AnalyticsRecorder* analytics_recorder_ = nullptr;
   std::string endpoint_id_ = "";
+  std::string local_endpoint_id_ = "";
 };
 
 }  // namespace nearby::connections

connections/implementation/base_pcp_handler.cc

@@ -894,13 +894,19 @@ ConnectionInfo BasePcpHandler::FillConnectionInfo(
     connection_info.ap_frequency = wifi_info.ap_frequency;
     if (NearbyFlags::GetInstance().GetBoolFlag(
             config_package_nearby::nearby_connections_feature::
-                kEnableDynamicRoleSwitch) &&
-        client->GetLocalOsInfo().type() == OsInfo::APPLE) {
-      ::location::nearby::connections::MediumRole medium_role_info;
-      medium_role_info.set_support_awdl_publisher(true);
-      medium_role_info.set_support_awdl_subscriber(true);
-      medium_role_info.set_support_wifi_hotspot_client(true);
-      connection_info.medium_role.emplace(medium_role_info);
+                kEnableDynamicRoleSwitch)) {
+      LOG(INFO) << "kEnableDynamicRoleSwitch is enabled";
+      ClientProxy::MediumsAvailability mediums_availability;
+      mediums_availability.is_wifi_direct_go_available =
+          mediums_->GetWifiDirect().IsGOAvailable();
+      mediums_availability.is_wifi_direct_gc_available =
+          mediums_->GetWifiDirect().IsGCAvailable();
+      mediums_availability.is_wifi_hotspot_ap_available =
+          mediums_->GetWifiHotspot().IsAPAvailable();
+      mediums_availability.is_wifi_hotspot_client_available =
+          mediums_->GetWifiHotspot().IsClientAvailable();
+      connection_info.medium_role.emplace(
+          client->GetLocalMediumRole(mediums_availability));
     }
     LOG(INFO) << "Query for WIFI information: is_supports_5_ghz="
               << connection_info.supports_5_ghz

connections/implementation/base_pcp_handler_test.cc

@@ -2949,5 +2949,37 @@ TEST_F(BasePcpHandlerTest, TestForceUpdateEndpointIdAdvertisingOption) {
   env_.Stop();
 }
 
+TEST_P(BasePcpHandlerTest,
+       FillConnectionInfo_kEnableDynamicRoleSwitch_Enabled) {
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kEnableDynamicRoleSwitch,
+      true);
+
+  Mediums m;
+  EndpointChannelManager ecm;
+  EndpointManager em(&ecm);
+  BwuManager bwu(m, em, ecm, {}, {});
+  MockPcpHandler pcp_handler(&m, &em, &ecm, &bwu);
+
+  ConnectionRequestInfo request_info = {
+      .endpoint_info = ByteArray("EndpointInfo"),
+  };
+  ConnectionOptions connection_options = {};
+
+  // Call FillConnectionInfo with dynamic role switch enabled
+  ConnectionInfo connection_info = pcp_handler.FillConnectionInfo(
+      client_.get(), request_info, connection_options);
+
+  // Verify that medium_role is set (populated) in connection_info!
+  EXPECT_TRUE(connection_info.medium_role.has_value());
+
+  bwu.Shutdown();
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kEnableDynamicRoleSwitch,
+      false);
+}
+
 }  // namespace
 }  // namespace nearby::connections

connections/implementation/bwu_manager.cc

@@ -297,9 +297,13 @@ void BwuManager::InitiateBwuForEndpoint(ClientProxy* client,
     if (is_dynamic_role_switch_enabled_ &&
         client->GetMediumRole(endpoint_id).has_value()) {
       MediumRole medium_role = client->GetMediumRole(endpoint_id).value();
-      if (NeedToSwitchRole(client, endpoint_id, proposed_medium, medium_role)) {
+      auto remote_os_info = client->GetRemoteOsInfo(endpoint_id);
+
+      if (NeedToSwitchRole(client, endpoint_id, proposed_medium, medium_role,
+                           remote_os_info.value_or(OsInfo()))) {
         if (!channel
                  ->Write(parser::ForBwuPathRequest(
+                     proposed_medium,
                      client->GetUpgradeMediums(endpoint_id).GetMediums(true),
                      medium_role))
                  .Ok()) {
@@ -571,6 +575,12 @@ void BwuManager::OnBwuNegotiationFrame(
           /* record_analytic= */ true,
           OperationResultCode::NEARBY_GENERIC_REMOTE_UPGRADE_FAILURE);
       break;
+    case BandwidthUpgradeNegotiationFrame::UPGRADE_PATH_REQUEST:
+      if (frame.upgrade_path_info().has_upgrade_path_request()) {
+        ProcessUpgradePathRequest(client, endpoint_id,
+                                  frame.upgrade_path_info());
+      }
+      break;
     case BandwidthUpgradeNegotiationFrame::LAST_WRITE_TO_PRIOR_CHANNEL:
       if (!in_progress_upgrades_.contains(endpoint_id)) {
         LOG(ERROR) << "Received LAST_WRITE_TO_PRIOR_CHANNEL for endpoint "
@@ -656,7 +666,19 @@ void BwuManager::OnIncomingConnection(
                "OfflineFrame on EndpointChannel "
             << channel->GetName();
 
-    const std::string& endpoint_id = introduction.endpoint_id();
+    std::string endpoint_id = introduction.endpoint_id();
+    if (is_dynamic_role_switch_enabled_ &&
+        !in_progress_upgrades_.contains(endpoint_id) &&
+        introduction.has_last_endpoint_id() &&
+        !introduction.last_endpoint_id().empty()) {
+      std::string last_endpoint_id = introduction.last_endpoint_id();
+      if (in_progress_upgrades_.contains(last_endpoint_id)) {
+        LOG(INFO) << "BwuManager: aliasing endpoint ID " << endpoint_id
+                  << " to " << last_endpoint_id;
+        endpoint_id = last_endpoint_id;
+      }
+    }
+
     ClientProxy* mapped_client;
     const auto item = in_progress_upgrades_.find(endpoint_id);
     if (item == in_progress_upgrades_.end()) return;
@@ -710,6 +732,7 @@ void BwuManager::RunOnBwuManagerThread(const std::string& name,
 void BwuManager::RunUpgradeProtocol(
     ClientProxy* client, const std::string& endpoint_id,
     std::unique_ptr<EndpointChannel> new_channel, bool enable_encryption) {
+  new_channel->SetLocalEndpointId(client->GetLocalEndpointId());
   LOG(INFO) << "RunUpgradeProtocol new channel @" << new_channel.get()
             << " name: " << new_channel->GetName() << ", medium: "
             << location::nearby::proto::connections::Medium_Name(
@@ -804,9 +827,11 @@ void BwuManager::ProcessBwuPathAvailableEvent(
       abort_bwu = true;
     } else {
       auto medium_role = client->GetMediumRole(endpoint_id);
+      auto remote_os_info = client->GetRemoteOsInfo(endpoint_id);
       if (medium_role.has_value() &&
           !NeedToSwitchRole(client, endpoint_id, upgrade_medium,
-                            medium_role.value())) {
+                            medium_role.value(),
+                            remote_os_info.value_or(OsInfo()))) {
         abort_bwu = true;
       }
     }
@@ -1013,9 +1038,20 @@ BwuManager::ProcessBwuPathAvailableEventInternal(
 
   // Write the requisite BANDWIDTH_UPGRADE_NEGOTIATION.CLIENT_INTRODUCTION as
   // the first OfflineFrame on this new EndpointChannel.
+  std::string last_local_endpoint_id = client->GetLastLocalEndpointId();
+  std::shared_ptr<EndpointChannel> previous_channel =
+      channel_manager_->GetChannelForEndpoint(endpoint_id);
+  if (previous_channel != nullptr) {
+    last_local_endpoint_id = previous_channel->GetLocalEndpointId();
+  }
+  LOG(INFO) << "BwuManager get last_local_endpoint_id "
+            << last_local_endpoint_id << " from "
+            << (previous_channel != nullptr ? "endpoint channel"
+                                            : "client proxy");
+
   if (!new_channel
            ->Write(parser::ForBwuIntroduction(
-               client->GetLocalEndpointId(),
+               client->GetLocalEndpointId(), last_local_endpoint_id,
                upgrade_path_info.supports_disabling_encryption()))
            .Ok()) {
     // This was never a fully EstablishedConnection, no need to provide a
@@ -1615,7 +1651,13 @@ void BwuManager::AttemptToRecordBandwidthUpgradeErrorForUnknownEndpoint(
 
 bool BwuManager::NeedToSwitchRole(
     ClientProxy* client, const std::string& endpoint_id, Medium medium,
-    const location::nearby::connections::MediumRole& medium_role) {
+    const location::nearby::connections::MediumRole& medium_role,
+    const location::nearby::connections::OsInfo& remote_os_info) {
+  if (!is_dynamic_role_switch_enabled_) {
+    return false;
+  }
+  // On called by receiver device, check if the sender device can host the
+  // upgrade medium or not
   if (GetLocalOsInfo(client).type() == OsInfo::APPLE) {
     switch (medium) {
       case Medium::WIFI_HOTSPOT:
@@ -1624,6 +1666,105 @@ bool BwuManager::NeedToSwitchRole(
         break;
     }
   }
+  // For testing on Windows as a receiver device to request dynamic role switch.
+  // No need for final check in.
+  if (GetLocalOsInfo(client).type() == OsInfo::WINDOWS &&
+     remote_os_info.type() == OsInfo::ANDROID) {
+    LOG(INFO) << "Local: Windows OS, Remote: Android device detected. "
+                 "WifiDirect NeedToSwitchRole and let Android be GO. "
+                 "medium_role.support_wifi_direct_group_owner(): "
+              << medium_role.support_wifi_direct_group_owner();
+    switch (medium) {
+      case Medium::WIFI_DIRECT:
+        return medium_role.support_wifi_direct_group_owner();
+      default:
+        break;
+    }
+  }
+
+  return false;
+}
+// This feature currently is only used by Android as receiver device to request
+// a dynamic role switch to Windows as Wi-Fi Direct GO. So WIFI_DIRECT is the
+// preferred medium to upgrade to.
+void BwuManager::ProcessUpgradePathRequest(
+    ClientProxy* client, const std::string& endpoint_id,
+    const location::nearby::connections::BandwidthUpgradeNegotiationFrame::
+        UpgradePathInfo& upgrade_path_info) {
+  if (!is_dynamic_role_switch_enabled_) {
+    return;
+  }
+  LOG(INFO) << "BwuManager: processing incoming UPGRADE_PATH_REQUEST frame for "
+               "endpoint "
+            << endpoint_id;
+
+  const auto& request = upgrade_path_info.upgrade_path_request();
+  std::vector<Medium> upgrade_mediums;
+  upgrade_mediums.reserve(request.mediums_size());
+  bool has_wifi_direct = false;
+  for (auto m : request.mediums()) {
+    Medium medium = parser::UpgradePathInfoMediumToMedium(
+        static_cast<BandwidthUpgradeNegotiationFrame::UpgradePathInfo::Medium>(
+            m));
+    LOG(INFO) << "BwuManager: UpgradePathRequest medium: "
+              << location::nearby::proto::connections::Medium_Name(medium);
+    upgrade_mediums.push_back(medium);
+    if (medium == Medium::WIFI_DIRECT) {
+      has_wifi_direct = true;
+    }
+  }
+
+  const location::nearby::connections::MediumRole& medium_role =
+      request.medium_meta_data().medium_role();
+  LOG(INFO) << "BwuManager: medium_role: " << medium_role.DebugString();
+
+  if (CanHost(client, medium_role)) {
+    Medium medium = ChooseBestUpgradeMedium(endpoint_id, upgrade_mediums);
+    if (has_wifi_direct) {
+      medium = Medium::WIFI_DIRECT;
+    }
+    LOG(INFO) << "BwuManager: Initiating BWU for endpoint " << endpoint_id
+              << " with medium "
+              << location::nearby::proto::connections::Medium_Name(medium);
+    InitiateBwuForEndpoint(client, endpoint_id, medium);
+  } else {
+    ProcessUpgradeFailureEvent(
+        client, endpoint_id, upgrade_path_info,
+        BandwidthUpgradeResult::REMOTE_CONNECTION_ERROR,
+        /* record_analytic= */ true,
+        OperationResultCode::NEARBY_GENERIC_REMOTE_UPGRADE_FAILURE);
+  }
+}
+
+bool BwuManager::CanHost(
+    ClientProxy* client,
+    const location::nearby::connections::MediumRole& medium_role) {
+  if (!is_dynamic_role_switch_enabled_) {
+    return false;
+  }
+  ClientProxy::MediumsAvailability mediums_availability;
+  mediums_availability.is_wifi_direct_go_available =
+      mediums_->GetWifiDirect().IsGOAvailable();
+  mediums_availability.is_wifi_direct_gc_available =
+      mediums_->GetWifiDirect().IsGCAvailable();
+  mediums_availability.is_wifi_hotspot_ap_available =
+      mediums_->GetWifiHotspot().IsAPAvailable();
+  mediums_availability.is_wifi_hotspot_client_available =
+      mediums_->GetWifiHotspot().IsClientAvailable();
+  const location::nearby::connections::MediumRole& local_medium_role =
+      client->GetLocalMediumRole(mediums_availability);
+  if ((local_medium_role.support_wifi_direct_group_owner() &&
+       medium_role.support_wifi_direct_group_client() &&
+       mediums_->GetWifiDirect().IsGOAvailable()) ||
+      (local_medium_role.support_wifi_hotspot_host() &&
+       medium_role.support_wifi_hotspot_client() &&
+       mediums_->GetWifiHotspot().IsAPAvailable()) ||
+      (local_medium_role.support_wifi_aware_publisher() &&
+       medium_role.support_wifi_aware_subscriber())) {
+    LOG(INFO) << "BwuManager: Can host the upgrade medium.";
+    return true;
+  }
+  LOG(INFO) << "BwuManager: Can't host the upgrade medium.";
   return false;
 }
 

connections/implementation/bwu_manager.h

@@ -221,7 +221,16 @@ class BwuManager : public EndpointManager::FrameProcessor {
 
   bool NeedToSwitchRole(
       ClientProxy* client, const std::string& endpoint_id, Medium medium,
-      const location::nearby::connections::MediumRole& medium_role);
+      const location::nearby::connections::MediumRole& medium_role,
+      const location::nearby::connections::OsInfo& remote_os_info);
+
+  void ProcessUpgradePathRequest(
+      ClientProxy* client, const std::string& endpoint_id,
+      const location::nearby::connections::BandwidthUpgradeNegotiationFrame::
+          UpgradePathInfo& upgrade_path_info);
+
+  bool CanHost(ClientProxy* client,
+               const location::nearby::connections::MediumRole& medium_role);
 
   virtual const location::nearby::connections::OsInfo& GetLocalOsInfo(
       ClientProxy* client) const;

connections/implementation/bwu_manager_test.cc

@@ -21,6 +21,7 @@
 #include "gtest/gtest.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/strings/string_view.h"
+#include "absl/time/time.h"
 #include "connections/connection_options.h"
 #include "connections/implementation/analytics/analytics_recorder.h"
 #include "connections/implementation/bwu_handler.h"
@@ -83,10 +84,13 @@ CreateWifiHotspotCredentials() {
 
 class BwuManagerTest : public ::testing::Test {
  protected:
-  BwuManagerTest() {
+  static void SetUpTestSuite() {
     NearbyFlags::GetInstance().OverrideBoolFlagValue(
         config_package_nearby::nearby_connections_feature::kEnableWifiDirect,
         true);
+  }
+
+  BwuManagerTest() {
     // Set up fake BWU handlers for WebRTC and WifiLAN.
     absl::flat_hash_map<Medium, std::unique_ptr<BwuHandler>> handlers;
     auto fake_web_rtc = std::make_unique<FakeBwuHandler>(Medium::WEB_RTC);
@@ -316,6 +320,141 @@ TEST(BwuManagerBaseTest, InitiateBwu_NeedToSwitchRole_Success) {
       false);
 }
 
+TEST(BwuManagerBaseTest, InitiateBwu_NeedToSwitchRole_WindowsAndroid_Success) {
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kEnableDynamicRoleSwitch,
+      true);
+  ClientProxy client;
+  EndpointChannelManager ecm;
+  EndpointManager em(&ecm);
+  Mediums mediums;
+  BwuManager::Config config;
+  config.allow_upgrade_to.SetAll(false);
+  absl::flat_hash_map<Medium, std::unique_ptr<BwuHandler>> handlers;
+  auto bwu_manager = std::make_unique<BwuManager>(mediums, em, ecm,
+                                                  std::move(handlers), config);
+
+  // Set up local as WINDOWS, remote as ANDROID
+  client.SetLocalOsType(OsInfo::WINDOWS);
+  OsInfo remote_os_info;
+  remote_os_info.set_type(OsInfo::ANDROID);
+
+  auto channel1 = std::make_unique<FakeEndpointChannel>(
+      Medium::BLUETOOTH, std::string(kServiceIdA));
+  auto* channel1_ptr = channel1.get();
+
+  MediumRole remote_medium_role;
+  remote_medium_role.set_support_wifi_direct_group_owner(true);
+
+  client.OnConnectionInitiated(
+      std::string(kEndpointId1),
+      {.remote_endpoint_info = ByteArray("remote endpoint")},
+      {.auto_upgrade_bandwidth = false,
+       .connection_info =
+           {
+               .medium_role = {remote_medium_role},
+           }},
+      {}, "");
+  client.OnConnectionAccepted(std::string(kEndpointId1));
+  client.SetRemoteOsInfo(kEndpointId1, remote_os_info);
+
+  ecm.RegisterChannelForEndpoint(&client, std::string(kEndpointId1),
+                                 std::move(channel1));
+
+  // Verify that before upgrade, write_timestamp is infinite past
+  EXPECT_EQ(channel1_ptr->GetLastWriteTimestamp(), absl::InfinitePast());
+
+  // Initiate BWU for WiFi Direct on Windows, which forces
+  // role switch to Android
+  bwu_manager->InitiateBwuForEndpoint(&client, std::string(kEndpointId1),
+                                      Medium::WIFI_DIRECT);
+
+  // Since role is switched, upgrade is NOT initiated locally, but delegated
+  EXPECT_FALSE(bwu_manager->IsUpgradeOngoing(std::string(kEndpointId1)));
+
+  // Verify that an UPGRADE_PATH_REQUEST frame was actually written to
+  // the channel
+  EXPECT_NE(channel1_ptr->GetLastWriteTimestamp(), absl::InfinitePast());
+
+  ecm.UnregisterChannelForEndpoint(std::string(kEndpointId1),
+                                   DisconnectionReason::LOCAL_DISCONNECTION,
+                                   SafeDisconnectionResult::kSafeDisconnection);
+  bwu_manager->Shutdown();
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kEnableDynamicRoleSwitch,
+      false);
+}
+
+TEST(BwuManagerBaseTest,
+     InitiateBwu_NeedToSwitchRole_WindowsAndroid_NoSwitch_Success) {
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kEnableDynamicRoleSwitch,
+      true);
+  ClientProxy client;
+  EndpointChannelManager ecm;
+  EndpointManager em(&ecm);
+  Mediums mediums;
+  BwuManager::Config config;
+  config.allow_upgrade_to.SetAll(false);
+  config.allow_upgrade_to.wifi_direct = true;
+  absl::flat_hash_map<Medium, std::unique_ptr<BwuHandler>> handlers;
+  auto fake_wifi_direct = std::make_unique<FakeBwuHandler>(Medium::WIFI_DIRECT);
+  auto* fake_wifi_direct_ptr = fake_wifi_direct.get();
+  handlers.emplace(Medium::WIFI_DIRECT, std::move(fake_wifi_direct));
+  auto bwu_manager = std::make_unique<BwuManager>(mediums, em, ecm,
+                                                  std::move(handlers), config);
+
+  // Set up local as WINDOWS, remote as ANDROID
+  client.SetLocalOsType(OsInfo::WINDOWS);
+  OsInfo remote_os_info;
+  remote_os_info.set_type(OsInfo::ANDROID);
+
+  auto channel1 = std::make_unique<FakeEndpointChannel>(
+      Medium::BLUETOOTH, std::string(kServiceIdA));
+
+  MediumRole remote_medium_role;
+  remote_medium_role.set_support_wifi_direct_group_owner(false);
+
+  client.OnConnectionInitiated(
+      std::string(kEndpointId1),
+      {.remote_endpoint_info = ByteArray("remote endpoint")},
+      {.auto_upgrade_bandwidth = false,
+       .connection_info =
+           {
+               .medium_role = {remote_medium_role},
+           }},
+      {}, "");
+  client.OnConnectionAccepted(std::string(kEndpointId1));
+  client.SetRemoteOsInfo(kEndpointId1, remote_os_info);
+
+  ecm.RegisterChannelForEndpoint(&client, std::string(kEndpointId1),
+                                 std::move(channel1));
+
+  // Verify that before upgrade, upgrade is not ongoing
+  EXPECT_FALSE(bwu_manager->IsUpgradeOngoing(std::string(kEndpointId1)));
+
+  // Initiate BWU for WiFi Direct on Windows. Since remote doesn't support GO,
+  // we do not switch roles, so we host/upgrade locally.
+  bwu_manager->InitiateBwuForEndpoint(&client, std::string(kEndpointId1),
+                                      Medium::WIFI_DIRECT);
+
+  // Upgrade is ongoing locally
+  EXPECT_TRUE(bwu_manager->IsUpgradeOngoing(std::string(kEndpointId1)));
+  EXPECT_EQ(fake_wifi_direct_ptr->handle_initialize_calls().size(), 1u);
+
+  ecm.UnregisterChannelForEndpoint(std::string(kEndpointId1),
+                                   DisconnectionReason::LOCAL_DISCONNECTION,
+                                   SafeDisconnectionResult::kSafeDisconnection);
+  bwu_manager->Shutdown();
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kEnableDynamicRoleSwitch,
+      false);
+}
+
 class BwuManagerTestParam : public BwuManagerTest,
                             public ::testing::WithParamInterface<bool> {
  protected:
@@ -1174,6 +1313,210 @@ TEST_F(BwuManagerTest, ReceiveUnexpectedLastWriteBeforeUpgrade_NoWedge) {
   UnRegisterChannelForEndpoint(kEndpointId1);
 }
 
+TEST_F(BwuManagerTest, ProcessUpgradePathRequest_CanHost_True) {
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kEnableDynamicRoleSwitch,
+      true);
+
+  // Shutdown original BwuManager to clean up registrations cleanly.
+  bwu_manager_->Shutdown();
+
+  // Set up fake BWU handlers for WifiDirect.
+  absl::flat_hash_map<Medium, std::unique_ptr<BwuHandler>> handlers;
+  auto fake_wifi_direct =
+      std::make_unique<FakeBwuHandler>(Medium::WIFI_DIRECT);
+  FakeBwuHandler* fake_wifi_direct_handler_ptr = fake_wifi_direct.get();
+  handlers.emplace(Medium::WIFI_DIRECT, std::move(fake_wifi_direct));
+
+  BwuManager::Config config;
+  config.allow_upgrade_to = BooleanMediumSelector{.wifi_direct = true};
+
+  bwu_manager_ = std::make_unique<BwuManager>(
+      mediums_, em_, ecm_, std::move(handlers), config);
+  bwu_manager_->MakeSingleThreadedForTesting();
+
+  // Create initial connection
+  CreateInitialEndpoint(&client_, kServiceIdA, kEndpointId1, Medium::BLUETOOTH);
+
+  // Build the UpgradePathRequest frame using the parser helper
+  location::nearby::connections::MediumRole remote_medium_role;
+  remote_medium_role.set_support_wifi_direct_group_client(true);
+  std::string bytes = parser::ForBwuPathRequest(
+      Medium::WIFI_DIRECT, {Medium::WIFI_DIRECT}, remote_medium_role);
+  OfflineFrame frame;
+  frame.ParseFromString(bytes);
+
+  // Process the request
+  bwu_manager_->OnIncomingFrame(frame, std::string(kEndpointId1), &client_,
+                                Medium::BLUETOOTH);
+
+  // Verify that WiFi Direct BWU was initiated
+  EXPECT_EQ(fake_wifi_direct_handler_ptr->handle_initialize_calls().size(), 1u);
+
+  UnRegisterChannelForEndpoint(kEndpointId1);
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kEnableDynamicRoleSwitch,
+      false);
+}
+
+TEST_F(BwuManagerTest, ProcessUpgradePathRequest_CanHost_False) {
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kEnableDynamicRoleSwitch,
+      true);
+
+  // Shutdown original BwuManager to clean up registrations cleanly.
+  bwu_manager_->Shutdown();
+
+  // Set up fake BWU handlers for WifiDirect.
+  absl::flat_hash_map<Medium, std::unique_ptr<BwuHandler>> handlers;
+  auto fake_wifi_direct =
+      std::make_unique<FakeBwuHandler>(Medium::WIFI_DIRECT);
+  FakeBwuHandler* fake_wifi_direct_handler_ptr = fake_wifi_direct.get();
+  handlers.emplace(Medium::WIFI_DIRECT, std::move(fake_wifi_direct));
+
+  BwuManager::Config config;
+  config.allow_upgrade_to = BooleanMediumSelector{.wifi_direct = true};
+
+  bwu_manager_ = std::make_unique<BwuManager>(
+      mediums_, em_, ecm_, std::move(handlers), config);
+  bwu_manager_->MakeSingleThreadedForTesting();
+
+  // Create initial connection
+  CreateInitialEndpoint(&client_, kServiceIdA, kEndpointId1, Medium::BLUETOOTH);
+
+  // Build the UpgradePathRequest frame where remote doesn't support GC
+  location::nearby::connections::MediumRole remote_medium_role;
+  std::string bytes = parser::ForBwuPathRequest(
+      Medium::WIFI_DIRECT, {Medium::WIFI_DIRECT}, remote_medium_role);
+  OfflineFrame frame;
+  frame.ParseFromString(bytes);
+
+  // Process the request
+  bwu_manager_->OnIncomingFrame(frame, std::string(kEndpointId1), &client_,
+                                Medium::BLUETOOTH);
+
+  // Verify that WiFi Direct BWU was NOT initiated
+  EXPECT_EQ(fake_wifi_direct_handler_ptr->handle_initialize_calls().size(), 0u);
+
+  UnRegisterChannelForEndpoint(kEndpointId1);
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kEnableDynamicRoleSwitch,
+      false);
+}
+
+TEST_F(BwuManagerTest, ProcessUpgradePathRequest_DynamicRoleSwitchDisabled) {
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kEnableDynamicRoleSwitch,
+      false);
+
+  // Shutdown original BwuManager to clean up registrations cleanly.
+  bwu_manager_->Shutdown();
+
+  // Set up fake BWU handlers for WifiDirect.
+  absl::flat_hash_map<Medium, std::unique_ptr<BwuHandler>> handlers;
+  auto fake_wifi_direct =
+      std::make_unique<FakeBwuHandler>(Medium::WIFI_DIRECT);
+  FakeBwuHandler* fake_wifi_direct_handler_ptr = fake_wifi_direct.get();
+  handlers.emplace(Medium::WIFI_DIRECT, std::move(fake_wifi_direct));
+
+  BwuManager::Config config;
+  config.allow_upgrade_to = BooleanMediumSelector{.wifi_direct = true};
+
+  bwu_manager_ = std::make_unique<BwuManager>(
+      mediums_, em_, ecm_, std::move(handlers), config);
+  bwu_manager_->MakeSingleThreadedForTesting();
+
+  // Create initial connection
+  CreateInitialEndpoint(&client_, kServiceIdA, kEndpointId1, Medium::BLUETOOTH);
+
+  // Build the UpgradePathRequest frame where remote supports GC
+  location::nearby::connections::MediumRole remote_medium_role;
+  remote_medium_role.set_support_wifi_direct_group_client(true);
+  std::string bytes = parser::ForBwuPathRequest(
+      Medium::WIFI_DIRECT, {Medium::WIFI_DIRECT}, remote_medium_role);
+  OfflineFrame frame;
+  frame.ParseFromString(bytes);
+
+  // Process the request
+  bwu_manager_->OnIncomingFrame(frame, std::string(kEndpointId1), &client_,
+                                Medium::BLUETOOTH);
+
+  // Verify that WiFi Direct BWU was NOT initiated
+  EXPECT_EQ(fake_wifi_direct_handler_ptr->handle_initialize_calls().size(), 0u);
+
+  UnRegisterChannelForEndpoint(kEndpointId1);
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kEnableDynamicRoleSwitch,
+      false);
+}
+
+TEST_F(BwuManagerTest, OnIncomingConnection_EndpointAliasesToLastEndpointId) {
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kEnableDynamicRoleSwitch,
+      true);
+
+  // Shutdown original BwuManager to clean up registrations cleanly.
+  bwu_manager_->Shutdown();
+
+  // Create a new BwuManager with the overridden flag. We use WEB_RTC since
+  // it has a simple, standard BWU flow.
+  absl::flat_hash_map<Medium, std::unique_ptr<BwuHandler>> handlers;
+  auto fake_web_rtc = std::make_unique<FakeBwuHandler>(Medium::WEB_RTC);
+  handlers.emplace(Medium::WEB_RTC, std::move(fake_web_rtc));
+
+  BwuManager::Config config;
+  config.allow_upgrade_to = BooleanMediumSelector{.web_rtc = true};
+
+  bwu_manager_ = std::make_unique<BwuManager>(
+      mediums_, em_, ecm_, std::move(handlers), config);
+  bwu_manager_->MakeSingleThreadedForTesting();
+
+  // Create initial connection with the old endpoint ID "OldEndpoint"
+  CreateInitialEndpoint(&client_, kServiceIdA, "OldEndpoint",
+                        Medium::BLUETOOTH);
+
+  // Initiate upgrade for "OldEndpoint" (inserts into in_progress_upgrades_)
+  bwu_manager_->InitiateBwuForEndpoint(&client_, "OldEndpoint",
+                                       Medium::WEB_RTC);
+
+  // Now simulate incoming upgraded connection. Set introduction read output:
+  // - endpoint_id = "NewEndpoint"
+  // - last_endpoint_id = "OldEndpoint"
+  auto upgraded_channel = std::make_unique<FakeEndpointChannel>(
+      Medium::WEB_RTC, std::string(kServiceIdA));
+  FakeEndpointChannel* upgraded_channel_raw = upgraded_channel.get();
+
+  std::string intro_frame = parser::ForBwuIntroduction(
+      "NewEndpoint", "OldEndpoint", /*supports_disabling_encryption=*/false);
+  upgraded_channel->set_read_output(
+      ExceptionOr<ByteArray>(ByteArray(intro_frame)));
+
+  auto connection = std::make_unique<BwuHandler::IncomingSocketConnection>();
+  connection->channel = std::move(upgraded_channel);
+
+  // Invoke OnIncomingConnection
+  bwu_manager_->InvokeOnIncomingConnectionForTesting(&client_,
+                                                     std::move(connection));
+
+  // Verify that an Ack frame was written to upgraded_channel_raw
+  EXPECT_NE(upgraded_channel_raw->GetLastWriteTimestamp(),
+            absl::InfinitePast());
+
+  // Clean up
+  UnRegisterChannelForEndpoint("OldEndpoint");
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kEnableDynamicRoleSwitch,
+      false);
+}
+
 INSTANTIATE_TEST_SUITE_P(BwuManagerTestParam, BwuManagerTestParam,
                          testing::Bool());
 

connections/implementation/client_proxy.cc

@@ -376,9 +376,10 @@ void ClientProxy::SetBluetoothMacAddress(const std::string& endpoint_id,
 
 std::string ClientProxy::GenerateLocalEndpointId() {
   if (!cached_endpoint_id_.empty()) {
-    if (stable_endpoint_id_mode_) {
+    if (stable_endpoint_id_mode_ || HasOngoingConnection()) {
       LOG(INFO) << "ClientProxy [Local Endpoint Re-using cached "
-                    "endpoint id due to in stable endpoint id mode]: "
+                    "endpoint id due to in stable endpoint id mode or having "
+                    "ongoing connection]: "
                     "client="
                 << GetClientId()
                 << "; cached_endpoint_id_=" << cached_endpoint_id_;
@@ -875,6 +876,48 @@ bool ClientProxy::HasOngoingConnection() const {
          !GetConnectedEndpoints().empty();
 }
 
+bool ClientProxy::HasWifiDirectConnection() const {
+  MutexLock lock(&mutex_);
+  for (const auto& entry : connections_) {
+    if (entry.second.first.connected_medium == Medium::WIFI_DIRECT) {
+      LOG(INFO) << "ClientProxy [HasWifiDirectConnection]: true";
+      return true;
+    }
+  }
+  LOG(INFO) << "ClientProxy [HasWifiDirectConnection]: false";
+  return false;
+}
+
+bool ClientProxy::HasWifiHotspotConnection() const {
+  MutexLock lock(&mutex_);
+  for (const auto& entry : connections_) {
+    if (entry.second.first.connected_medium == Medium::WIFI_HOTSPOT) {
+      return true;
+    }
+  }
+  return false;
+}
+
+bool ClientProxy::HasWifiAwareConnection() const {
+  MutexLock lock(&mutex_);
+  for (const auto& entry : connections_) {
+    if (entry.second.first.connected_medium == Medium::WIFI_AWARE) {
+      return true;
+    }
+  }
+  return false;
+}
+
+std::string ClientProxy::GetLastLocalEndpointId() const {
+  MutexLock lock(&mutex_);
+  return last_local_endpoint_id_;
+}
+
+void ClientProxy::SetLastLocalEndpointId(absl::string_view endpoint_id) {
+  MutexLock lock(&mutex_);
+  last_local_endpoint_id_ = std::string(endpoint_id);
+}
+
 std::int32_t ClientProxy::GetNumOutgoingConnections() const {
   return GetMatchingEndpoints([](const Connection& connection) {
            return connection.status == Connection::kConnected &&
@@ -1249,9 +1292,23 @@ void ClientProxy::RemoveAllEndpoints() {
   OnSessionComplete();
 }
 
+void ClientProxy::ResetLocalEndpointId() {
+  MutexLock lock(&mutex_);
+  if (HasOngoingConnection()) {
+    return;
+  }
+  if (!local_endpoint_id_.empty()) {
+    last_local_endpoint_id_ = local_endpoint_id_;
+    local_endpoint_id_.clear();
+  }
+}
+
 void ClientProxy::OnSessionComplete() {
   MutexLock lock(&mutex_);
   if (connections_.empty() && !IsAdvertising()) {
+    if (!local_endpoint_id_.empty()) {
+      last_local_endpoint_id_ = local_endpoint_id_;
+    }
     local_endpoint_id_.clear();
 
     analytics_recorder_->LogSession();
@@ -1298,6 +1355,9 @@ void ClientProxy::EnterStableEndpointIdMode() {
           << GetClientId();
 
   stable_endpoint_id_mode_ = true;
+  if (!IsAdvertising() && !IsDiscovering() && !HasOngoingConnection()) {
+    ResetLocalEndpointId();
+  }
 }
 
 void ClientProxy::ExitStableEndpointIdMode() {
@@ -1305,6 +1365,7 @@ void ClientProxy::ExitStableEndpointIdMode() {
   VLOG(1) << "ClientProxy [ExitStableEndpointIdMode]: client=" << GetClientId();
 
   stable_endpoint_id_mode_ = false;
+  ResetLocalEndpointId();
   ScheduleClearCachedEndpointIdAlarm();
 }
 
@@ -1318,7 +1379,7 @@ void ClientProxy::ScheduleClearCachedEndpointIdAlarm() {
     return;
   }
 
-  if (HasOngoingConnection()) {
+  if (IsAdvertising() || IsDiscovering() || HasOngoingConnection()) {
     VLOG(1) << "ClientProxy [Handle clearing cached endpoint ID "
                "during disconnection]: client="
             << GetClientId();
@@ -1431,6 +1492,41 @@ std::optional<MediumRole> ClientProxy::GetMediumRole(
   return std::nullopt;
 }
 
+location::nearby::connections::MediumRole ClientProxy::GetLocalMediumRole(
+    const ClientProxy::MediumsAvailability& mediums_availability) const {
+  location::nearby::connections::MediumRole medium_role;
+  if (!NearbyFlags::GetInstance().GetBoolFlag(
+          config_package_nearby::nearby_connections_feature::
+              kEnableDynamicRoleSwitch)) {
+    return medium_role;
+  }
+
+  if (GetLocalOsInfo().type() == OsInfo::APPLE) {
+    medium_role.set_support_awdl_publisher(true);
+    medium_role.set_support_awdl_subscriber(true);
+    // Apple always supports wifi hotspot client role since they can always
+    // join a hotspot.
+    medium_role.set_support_wifi_hotspot_client(true);
+    return medium_role;
+  }
+
+  medium_role.set_support_wifi_direct_group_owner(
+      mediums_availability.is_wifi_direct_go_available && !IsUsingP2pMedium());
+  medium_role.set_support_wifi_direct_group_client(
+      mediums_availability.is_wifi_direct_gc_available);
+  medium_role.set_support_wifi_hotspot_host(
+      mediums_availability.is_wifi_hotspot_ap_available && !IsUsingP2pMedium());
+  medium_role.set_support_wifi_hotspot_client(
+      mediums_availability.is_wifi_hotspot_client_available);
+  LOG(INFO) << "medium_role: " << medium_role.DebugString();
+  return medium_role;
+}
+
+bool ClientProxy::IsUsingP2pMedium() const {
+  return HasWifiDirectConnection() || HasWifiHotspotConnection() ||
+         HasWifiAwareConnection();
+}
+
 std::optional<std::string> ClientProxy::GetEndpointIdForDct() const {
   MutexLock lock(&mutex_);
   if (dct_endpoint_id_.empty()) {

connections/implementation/client_proxy.h

@@ -103,6 +103,9 @@ class ClientProxy final {
   // Clears all the runtime state of this client.
   void Reset();
 
+  // Resets the local endpoint ID and sets the last local endpoint ID.
+  void ResetLocalEndpointId();
+
   // Marks this client as advertising with the given callbacks.
   void StartedAdvertising(
       const std::string& service_id, Strategy strategy,
@@ -210,6 +213,14 @@ class ClientProxy final {
   // Returns true if there is at least one connected connection or one pending
   // connection.
   bool HasOngoingConnection() const;
+  // Returns true if there is at least one active WiFi Direct connection.
+  bool HasWifiDirectConnection() const;
+  // Returns true if there is at least one active WiFi Hotspot connection.
+  bool HasWifiHotspotConnection() const;
+  // Returns true if there is at least one active WiFi Aware connection.
+  bool HasWifiAwareConnection() const;
+  std::string GetLastLocalEndpointId() const;
+  void SetLastLocalEndpointId(absl::string_view endpoint_id);
   // Returns the number of endpoints that are connected and outgoing.
   std::int32_t GetNumOutgoingConnections() const;
   // Returns the number of endpoints that are connected and incoming.
@@ -333,6 +344,18 @@ class ClientProxy final {
   std::optional<location::nearby::connections::MediumRole> GetMediumRole(
       absl::string_view endpoint_id) const;
 
+  struct MediumsAvailability {
+    bool is_wifi_direct_go_available = false;
+    bool is_wifi_direct_gc_available = false;
+    bool is_wifi_hotspot_ap_available = false;
+    bool is_wifi_hotspot_client_available = false;
+  };
+
+  location::nearby::connections::MediumRole GetLocalMediumRole(
+      const MediumsAvailability& mediums_availability) const;
+
+  bool IsUsingP2pMedium() const;
+
   // Forces client to regenerate a new local endpoint id.
   void ClearCachedLocalEndpointId();
 
@@ -441,6 +464,7 @@ class ClientProxy final {
   std::int64_t client_id_;
   std::string local_endpoint_id_;
   std::string local_endpoint_info_;
+  std::string last_local_endpoint_id_;
 
   // If advertising is in stable endpoint ID mode, the endpoint ID is stable
   // for 30s after advertising or disconnection. When stable_endpoint_id_mode_

connections/implementation/client_proxy_test.cc

@@ -1563,6 +1563,257 @@ TEST_F(ClientProxyTest, GetSavePathDefaultsToEmpty) {
   EXPECT_THAT(client1()->GetSavePath(advertising_endpoint.id), IsEmpty());
 }
 
+TEST_F(ClientProxyTest, GetLocalMediumRoleFlagDisabled) {
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kEnableDynamicRoleSwitch,
+      false);
+  ClientProxy::MediumsAvailability availability;
+  availability.is_wifi_direct_go_available = true;
+  availability.is_wifi_direct_gc_available = true;
+  availability.is_wifi_hotspot_ap_available = true;
+  availability.is_wifi_hotspot_client_available = true;
+
+  location::nearby::connections::MediumRole role =
+      client1()->GetLocalMediumRole(availability);
+  EXPECT_FALSE(role.support_awdl_publisher());
+  EXPECT_FALSE(role.support_awdl_subscriber());
+  EXPECT_FALSE(role.support_wifi_direct_group_owner());
+  EXPECT_FALSE(role.support_wifi_direct_group_client());
+  EXPECT_FALSE(role.support_wifi_hotspot_host());
+  EXPECT_FALSE(role.support_wifi_hotspot_client());
+}
+
+TEST_F(ClientProxyTest, GetLocalMediumRoleAppleOs) {
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kEnableDynamicRoleSwitch,
+      true);
+  client1()->SetLocalOsType(location::nearby::connections::OsInfo::APPLE);
+  ClientProxy::MediumsAvailability availability;
+
+  location::nearby::connections::MediumRole role =
+      client1()->GetLocalMediumRole(availability);
+  EXPECT_TRUE(role.support_awdl_publisher());
+  EXPECT_TRUE(role.support_awdl_subscriber());
+  EXPECT_TRUE(role.support_wifi_hotspot_client());
+  EXPECT_FALSE(role.support_wifi_direct_group_owner());
+  EXPECT_FALSE(role.support_wifi_direct_group_client());
+  EXPECT_FALSE(role.support_wifi_hotspot_host());
+}
+
+TEST_F(ClientProxyTest, GetLocalMediumRoleNonAppleOsNoP2pConnection) {
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kEnableDynamicRoleSwitch,
+      true);
+  client1()->SetLocalOsType(location::nearby::connections::OsInfo::ANDROID);
+  ClientProxy::MediumsAvailability availability;
+  availability.is_wifi_direct_go_available = true;
+  availability.is_wifi_direct_gc_available = true;
+  availability.is_wifi_hotspot_ap_available = true;
+  availability.is_wifi_hotspot_client_available = true;
+
+  location::nearby::connections::MediumRole role =
+      client1()->GetLocalMediumRole(availability);
+  EXPECT_TRUE(role.support_wifi_direct_group_owner());
+  EXPECT_TRUE(role.support_wifi_direct_group_client());
+  EXPECT_TRUE(role.support_wifi_hotspot_host());
+  EXPECT_TRUE(role.support_wifi_hotspot_client());
+}
+
+TEST_F(ClientProxyTest, GetLocalMediumRoleNonAppleOsWithP2pConnection) {
+  NearbyFlags::GetInstance().OverrideBoolFlagValue(
+      config_package_nearby::nearby_connections_feature::
+          kEnableDynamicRoleSwitch,
+      true);
+  client1()->SetLocalOsType(location::nearby::connections::OsInfo::ANDROID);
+
+  // Setup an active P2P connection to make IsUsingP2pMedium() true
+  Endpoint advertising_endpoint =
+      StartAdvertising(client1(), advertising_connection_listener_);
+  OnAdvertisingConnectionInitiated(client1(), advertising_endpoint);
+  client1()->OnBandwidthChanged(advertising_endpoint.id, Medium::WIFI_DIRECT);
+  EXPECT_TRUE(client1()->IsUsingP2pMedium());
+
+  ClientProxy::MediumsAvailability availability;
+  availability.is_wifi_direct_go_available = true;
+  availability.is_wifi_direct_gc_available = true;
+  availability.is_wifi_hotspot_ap_available = true;
+  availability.is_wifi_hotspot_client_available = true;
+
+  location::nearby::connections::MediumRole role =
+      client1()->GetLocalMediumRole(availability);
+  EXPECT_FALSE(role.support_wifi_direct_group_owner());
+  EXPECT_TRUE(role.support_wifi_direct_group_client());
+  EXPECT_FALSE(role.support_wifi_hotspot_host());
+  EXPECT_TRUE(role.support_wifi_hotspot_client());
+}
+
+TEST_F(ClientProxyTest, GetNumIncomingAndOutgoingConnections) {
+  // Initially no connections
+  EXPECT_EQ(client1()->GetNumIncomingConnections(), 0);
+  EXPECT_EQ(client1()->GetNumOutgoingConnections(), 0);
+
+  // Set expectation for acceptance callback on step 1
+  // (which is outgoing based on discovery_connection_info_)
+  EXPECT_CALL(mock_advertising_connection_.accepted_cb, Call).Times(1);
+
+  // Define a complete listener for advertising
+  ConnectionListener advertising_listener = {
+      .initiated_cb = mock_advertising_connection_.initiated_cb.AsStdFunction(),
+      .accepted_cb = mock_advertising_connection_.accepted_cb.AsStdFunction(),
+  };
+
+  // 1. Establish connection 1
+  Endpoint advertising_endpoint =
+      StartAdvertising(client1(), advertising_listener);
+  EXPECT_CALL(mock_advertising_connection_.initiated_cb, Call).Times(1);
+  client1()->OnConnectionInitiated(
+      advertising_endpoint.id, discovery_connection_info_, connection_options_,
+      advertising_listener, "connection_token1");
+
+  // Accept local, accept remote, and then OnConnectionAccepted
+  client1()->LocalEndpointAcceptedConnection(
+      advertising_endpoint.id,
+      {
+          .payload_cb = mock_discovery_payload_.payload_cb.AsStdFunction(),
+          .payload_progress_cb =
+              mock_discovery_payload_.payload_progress_cb.AsStdFunction(),
+      });
+  client1()->RemoteEndpointAcceptedConnection(advertising_endpoint.id);
+  client1()->OnConnectionAccepted(advertising_endpoint.id);
+
+  // Verify client1 has 0 incoming connections and 1 outgoing connection
+  EXPECT_EQ(client1()->GetNumIncomingConnections(), 0);
+  EXPECT_EQ(client1()->GetNumOutgoingConnections(), 1);
+
+  // Set expectation for acceptance callback on step 2
+  // (which is incoming based on advertising_connection_info_)
+  EXPECT_CALL(mock_discovery_connection_.accepted_cb, Call).Times(1);
+
+  // 2. Establish connection 2
+  StartDiscovery(client1(), GetDiscoveryListener());
+  Endpoint remote_endpoint = {
+      .info = ByteArray{"remote endpoint name"},
+      .id = "rem_ep_id",
+  };
+  OnDiscoveryEndpointFound(client1(), remote_endpoint);
+
+  EXPECT_CALL(mock_discovery_connection_.initiated_cb, Call).Times(1);
+  client1()->OnConnectionInitiated(
+      remote_endpoint.id, advertising_connection_info_, connection_options_,
+      discovery_connection_listener_, "connection_token2");
+
+  // Accept local, accept remote, and then OnConnectionAccepted
+  client1()->LocalEndpointAcceptedConnection(
+      remote_endpoint.id,
+      {
+          .payload_cb = mock_discovery_payload_.payload_cb.AsStdFunction(),
+          .payload_progress_cb =
+              mock_discovery_payload_.payload_progress_cb.AsStdFunction(),
+      });
+  client1()->RemoteEndpointAcceptedConnection(remote_endpoint.id);
+  client1()->OnConnectionAccepted(remote_endpoint.id);
+
+  // Verify client1 has 1 incoming connection and 1 outgoing connection
+  EXPECT_EQ(client1()->GetNumIncomingConnections(), 1);
+  EXPECT_EQ(client1()->GetNumOutgoingConnections(), 1);
+}
+
+TEST_F(ClientProxyTest, IsUsingP2pMediumTests) {
+  // With no connections, IsUsingP2pMedium should be false
+  EXPECT_FALSE(client1()->IsUsingP2pMedium());
+
+  // 1. Connection with Non-P2P medium (e.g. WIFI_LAN)
+  Endpoint endpoint_lan =
+      StartAdvertising(client1(), advertising_connection_listener_);
+  OnAdvertisingConnectionInitiated(client1(), endpoint_lan);
+  client1()->OnBandwidthChanged(endpoint_lan.id, Medium::WIFI_LAN);
+  EXPECT_FALSE(client1()->IsUsingP2pMedium());
+
+  // Clean-up connection
+  client1()->OnDisconnected(endpoint_lan.id, /*notify=*/false);
+  EXPECT_FALSE(client1()->IsUsingP2pMedium());
+
+  // 2. Connection with WIFI_DIRECT
+  Endpoint endpoint_direct =
+      StartAdvertising(client1(), advertising_connection_listener_);
+  OnAdvertisingConnectionInitiated(client1(), endpoint_direct);
+  client1()->OnBandwidthChanged(endpoint_direct.id, Medium::WIFI_DIRECT);
+  EXPECT_TRUE(client1()->IsUsingP2pMedium());
+  client1()->OnDisconnected(endpoint_direct.id, /*notify=*/false);
+
+  // 3. Connection with WIFI_HOTSPOT
+  Endpoint endpoint_hotspot =
+      StartAdvertising(client1(), advertising_connection_listener_);
+  OnAdvertisingConnectionInitiated(client1(), endpoint_hotspot);
+  client1()->OnBandwidthChanged(endpoint_hotspot.id, Medium::WIFI_HOTSPOT);
+  EXPECT_TRUE(client1()->IsUsingP2pMedium());
+  client1()->OnDisconnected(endpoint_hotspot.id, /*notify=*/false);
+
+  // 4. Connection with WIFI_AWARE
+  Endpoint endpoint_aware =
+      StartAdvertising(client1(), advertising_connection_listener_);
+  OnAdvertisingConnectionInitiated(client1(), endpoint_aware);
+  client1()->OnBandwidthChanged(endpoint_aware.id, Medium::WIFI_AWARE);
+  EXPECT_TRUE(client1()->IsUsingP2pMedium());
+  client1()->OnDisconnected(endpoint_aware.id, /*notify=*/false);
+  EXPECT_FALSE(client1()->IsUsingP2pMedium());
+}
+
+TEST_F(ClientProxyTest, GetAndSetLastLocalEndpointId) {
+  EXPECT_TRUE(client1()->GetLastLocalEndpointId().empty());
+  client1()->SetLastLocalEndpointId("TestEndpointID");
+  EXPECT_EQ(client1()->GetLastLocalEndpointId(), "TestEndpointID");
+}
+
+TEST_F(ClientProxyTest, ResetLocalEndpointId_OngoingConnectionReturnsEarly) {
+  std::string old_id = client1()->GetLocalEndpointId();
+  ASSERT_FALSE(old_id.empty());
+
+  // Set up an ongoing connection
+  OnAdvertisingConnectionInitiated(client1(),
+                                    {ByteArray("EndpointInfo"), "EndA"});
+  EXPECT_TRUE(client1()->HasOngoingConnection());
+
+  // ResetLocalEndpointId should NOT clear local_endpoint_id
+  client1()->ResetLocalEndpointId();
+  EXPECT_EQ(client1()->GetLocalEndpointId(), old_id);
+
+  // Terminate connection
+  client1()->OnDisconnected("EndA", /*notify=*/false);
+  EXPECT_FALSE(client1()->HasOngoingConnection());
+
+  // ResetLocalEndpointId should now successfully clear local_endpoint_id
+  client1()->ResetLocalEndpointId();
+  EXPECT_NE(client1()->GetLocalEndpointId(), old_id);
+}
+
+TEST_F(ClientProxyTest, ResetLocalEndpointId_SavesToLastLocalEndpointId) {
+  std::string old_id = client1()->GetLocalEndpointId();
+  ASSERT_FALSE(old_id.empty());
+
+  client1()->ResetLocalEndpointId();
+  EXPECT_EQ(client1()->GetLastLocalEndpointId(), old_id);
+}
+
+TEST_F(ClientProxyTest, OnSessionComplete_SavesToLastLocalEndpointId) {
+  std::string old_id = client1()->GetLocalEndpointId();
+  ASSERT_FALSE(old_id.empty());
+
+  // Put client into advertising mode first
+  client1()->StartedAdvertising(service_id_, strategy_, {}, {}, {});
+  EXPECT_TRUE(client1()->IsAdvertising());
+
+  // Stopping advertising triggers OnSessionComplete.
+  // Since connections_ is empty, it completes the session and should save last
+  // endpoint ID.
+  client1()->StoppedAdvertising();
+  EXPECT_FALSE(client1()->IsAdvertising());
+  EXPECT_EQ(client1()->GetLastLocalEndpointId(), old_id);
+}
+
 }  // namespace
 }  // namespace connections
 }  // namespace nearby

connections/implementation/encryption_runner_test.cc

@@ -103,6 +103,10 @@ class FakeEndpointChannel : public EndpointChannel {
   }
   void SetAnalyticsRecorder(analytics::AnalyticsRecorder* analytics_recorder,
                             const std::string& endpoint_id) override {}
+  void SetLocalEndpointId(const std::string& local_endpoint_id) override {
+    local_endpoint_id_ = local_endpoint_id;
+  }
+  std::string GetLocalEndpointId() const override { return local_endpoint_id_; }
 
  private:
   InputStream* in_ = nullptr;
@@ -110,6 +114,7 @@ class FakeEndpointChannel : public EndpointChannel {
   absl::Time read_timestamp_ = absl::InfinitePast();
   absl::Time write_timestamp_ = absl::InfinitePast();
   mutable uint32_t next_keep_alive_seq_no_ = 0;
+  std::string local_endpoint_id_;
 };
 
 struct User {

connections/implementation/endpoint_channel.h

@@ -128,6 +128,9 @@ class EndpointChannel {
 
   // Enables the multiplex socket on the EndpointChannel.
   virtual bool EnableMultiplexSocket() { return false; }
+
+  virtual void SetLocalEndpointId(const std::string& local_endpoint_id) = 0;
+  virtual std::string GetLocalEndpointId() const = 0;
 };
 
 inline bool operator==(const EndpointChannel& lhs, const EndpointChannel& rhs) {

connections/implementation/endpoint_channel_manager.cc

@@ -107,6 +107,7 @@ void EndpointChannelManager::SetActiveEndpointChannel(
   // Update the channel first, then encrypt this new channel, if
   // crypto context is present.
   channel->SetAnalyticsRecorder(&client->GetAnalyticsRecorder(), endpoint_id);
+  channel->SetLocalEndpointId(client->GetLocalEndpointId());
   channel_state_.UpdateChannelForEndpoint(endpoint_id, std::move(channel));
   channel_state_.UpdateSafeToDisconnectForEndpoint(
       endpoint_id, client->IsSafeToDisconnectEnabled(endpoint_id));

connections/implementation/fake_bwu_handler.h

@@ -89,6 +89,7 @@ class FakeBwuHandler : public BaseBwuHandler {
     upgraded_channel->set_read_output(
         ExceptionOr<ByteArray>(ByteArray(parser::ForBwuIntroduction(
             *handle_initialize_calls_[initialize_call_index].endpoint_id,
+            /*last_endpoint_id=*/"",
             false /* supports_disabling_encryption */))));
     auto connection = std::make_unique<IncomingSocketConnection>();
     connection->channel = std::move(upgraded_channel);

connections/implementation/fake_endpoint_channel.h

@@ -93,6 +93,10 @@ class FakeEndpointChannel : public EndpointChannel {
   }
   void SetAnalyticsRecorder(analytics::AnalyticsRecorder* analytics_recorder,
                             const std::string& endpoint_id) override {}
+  void SetLocalEndpointId(const std::string& local_endpoint_id) override {
+    local_endpoint_id_ = local_endpoint_id;
+  }
+  std::string GetLocalEndpointId() const override { return local_endpoint_id_; }
 
   void set_read_output(ExceptionOr<ByteArray> output) { read_output_ = output; }
   void set_write_output(Exception output) { write_output_ = output; }
@@ -113,6 +117,7 @@ class FakeEndpointChannel : public EndpointChannel {
   bool is_paused_ = false;
   location::nearby::proto::connections::DisconnectionReason
       disconnection_reason_;
+  std::string local_endpoint_id_;
   mutable uint32_t next_keep_alive_seq_no_ = 0;
 };
 

connections/implementation/mock_endpoint_channel.h

@@ -69,6 +69,9 @@ class MockEndpointChannel : public EndpointChannel {
   MOCK_METHOD(uint32_t, GetNextKeepAliveSeqNo, (), (const, override));
   MOCK_METHOD(void, SetAnalyticsRecorder,
               (analytics::AnalyticsRecorder*, const std::string&), (override));
+  MOCK_METHOD(void, SetLocalEndpointId, (const std::string& local_endpoint_id),
+              (override));
+  MOCK_METHOD(std::string, GetLocalEndpointId, (), (const, override));
 };
 
 }  // namespace nearby::connections

connections/implementation/offline_frames.cc

@@ -460,6 +460,7 @@ std::string ForBwuSafeToClose() {
 }
 
 std::string ForBwuIntroduction(const std::string& endpoint_id,
+                             const std::string& last_endpoint_id,
                              bool supports_disabling_encryption) {
   OfflineFrame frame;
 
@@ -473,6 +474,9 @@ std::string ForBwuIntroduction(const std::string& endpoint_id,
   client_introduction->set_endpoint_id(endpoint_id);
   client_introduction->set_supports_disabling_encryption(
       supports_disabling_encryption);
+  if (!last_endpoint_id.empty()) {
+    client_introduction->set_last_endpoint_id(last_endpoint_id);
+  }
 
   return frame.SerializeAsString();
 }
@@ -498,15 +502,12 @@ std::string ForBwuFailure(const UpgradePathInfo& info) {
   v1_frame->set_type(V1Frame::BANDWIDTH_UPGRADE_NEGOTIATION);
   auto* sub_frame = v1_frame->mutable_bandwidth_upgrade_negotiation();
   sub_frame->set_event_type(BandwidthUpgradeNegotiationFrame::UPGRADE_FAILURE);
-  auto* upgrade_path_info = sub_frame->mutable_upgrade_path_info();
-  *upgrade_path_info = info;
-
   *sub_frame->mutable_upgrade_path_info() = info;
 
   return frame.SerializeAsString();
 }
 
-std::string ForBwuPathRequest(const std::vector<Medium>& mediums,
+std::string ForBwuPathRequest(Medium medium, const std::vector<Medium>& mediums,
                             const MediumRole& medium_role) {
   OfflineFrame frame;
 
@@ -516,8 +517,10 @@ std::string ForBwuPathRequest(const std::vector<Medium>& mediums,
   auto* sub_frame = v1_frame->mutable_bandwidth_upgrade_negotiation();
   sub_frame->set_event_type(
       BandwidthUpgradeNegotiationFrame::UPGRADE_PATH_REQUEST);
+  auto* upgrade_path_info = sub_frame->mutable_upgrade_path_info();
+  upgrade_path_info->set_medium(MediumToUpgradePathInfoMedium(medium));
   auto* upgrade_path_request =
-      sub_frame->mutable_upgrade_path_info()->mutable_upgrade_path_request();
+      upgrade_path_info->mutable_upgrade_path_request();
   for (const auto& medium : mediums) {
     upgrade_path_request->add_mediums(MediumToUpgradePathInfoMedium(medium));
   }

connections/implementation/offline_frames.h

@@ -76,6 +76,7 @@ std::string ForPayloadAckPayloadTransfer(std::int64_t payload_id);
 
 // Builds Bandwidth Upgrade [BWU] messages.
 std::string ForBwuIntroduction(const std::string& endpoint_id,
+                             const std::string& last_endpoint_id,
                              bool supports_disabling_encryption);
 std::string ForBwuIntroductionAck();
 std::string ForBwuWifiHotspotPathAvailable(
@@ -107,7 +108,7 @@ std::string ForBwuWebrtcPathAvailable(
     const location::nearby::connections::LocationHint& location_hint_a);
 std::string ForBwuFailure(const UpgradePathInfo& info);
 std::string ForBwuPathRequest(
-    const std::vector<Medium>& mediums,
+    Medium medium, const std::vector<Medium>& mediums,
     const location::nearby::connections::MediumRole& medium_role);
 std::string ForBwuLastWrite();
 std::string ForBwuSafeToClose();

connections/implementation/offline_frames_test.cc

@@ -502,10 +502,7 @@ TEST(OfflineFramesTest, CanGenerateBwuWifiLanPathAvailable) {
               ip_address: "\x2a\x00\x79\xe0\x2e\x87\x00\x06\xb7\x28\x67\x45\x7a\xdd\x01\x53"
               port: 1234
             >
-            address_candidates: <
-              ip_address: "\001\002\003\004"
-              port: 1234
-            >
+            address_candidates: < ip_address: "\001\002\003\004" port: 1234 >
           >
           supports_client_introduction_ack: true
         >
@@ -679,11 +676,13 @@ TEST(OfflineFramesTest, CanGenerateBwuIntroduction) {
         client_introduction: <
           endpoint_id: "ABC"
           supports_disabling_encryption: false
+          last_endpoint_id: "DEF"
         >
       >
     >)pb";
-  auto response = FromBytes(ForBwuIntroduction(
-      std::string(kEndpointId), false /* supports_disabling_encryption */));
+  auto response =
+      FromBytes(ForBwuIntroduction(std::string(kEndpointId), "DEF",
+                                   false /* supports_disabling_encryption */));
   ASSERT_TRUE(response.ok());
   OfflineFrame message = response.result();
   EXPECT_THAT(message, EqualsProto(kExpected));
@@ -722,7 +721,6 @@ TEST(OfflineFramesTest, CanGenerateDisconnection) {
   EXPECT_THAT(message, EqualsProto(kExpected));
 }
 
-
 TEST(OfflineFramesTest, CanGenerateBwuPathRequest) {
   constexpr absl::string_view kExpected =
       R"pb(
@@ -732,6 +730,7 @@ TEST(OfflineFramesTest, CanGenerateBwuPathRequest) {
       bandwidth_upgrade_negotiation: <
         event_type: UPGRADE_PATH_REQUEST
         upgrade_path_info: <
+          medium: WIFI_HOTSPOT
           upgrade_path_request: <
             mediums: WIFI_HOTSPOT
             medium_meta_data: <
@@ -745,7 +744,8 @@ TEST(OfflineFramesTest, CanGenerateBwuPathRequest) {
   mediums.push_back(Medium::WIFI_HOTSPOT);
   MediumRole medium_role;
   medium_role.set_support_wifi_hotspot_client(true);
-  auto response = FromBytes(ForBwuPathRequest(mediums, medium_role));
+  auto response =
+      FromBytes(ForBwuPathRequest(Medium::WIFI_HOTSPOT, mediums, medium_role));
   ASSERT_TRUE(response.ok());
   OfflineFrame message = response.result();
   EXPECT_THAT(message, EqualsProto(kExpected));

connections/implementation/offline_frames_validator_test.cc

@@ -765,7 +765,7 @@ TEST(OfflineFramesValidatorTest,
   OfflineFrame offline_frame_2;
 
   std::string wifi_direct_ssid{"DIRECT-A*-0123456789AB"};
-  std::string wifi_direct_pin_wrong_length = "abcefghijklmnopqrstuvwxyz";
+  std::string wifi_direct_pin_wrong_length = "01234567890123456";
   std::string bytes = ForBwuWifiDirectPathAvailable(
       wifi_direct_ssid, std::string(kWifiDirectPassword), kPort,
       kWifiDirectFrequency, kSupportsDisablingEncryption,