Snap for 9922117 from dba24d4531b23ab415400dfec628894ae7d00ebd to studio-giraffe-releasestudio-2022.3.1-rc1studio-2022.3.1-beta2studio-2022.3.1studio-main

Change-Id: I2e42ea2bae2b28571a071f111e2b4e67fb90fa95

15 files changed, 633 insertions, 428 deletions

diff --git a/directaccess/src/META-INF/plugin.xml b/directaccess/src/META-INF/plugin.xml
index acd65d1..10f8b48 100644
--- a/directaccess/src/META-INF/plugin.xml
+++ b/directaccess/src/META-INF/plugin.xml
@@ -32,6 +32,6 @@
   </extensions>
 
   <extensions defaultExtensionNs="com.android.tools.idea">
-    <deviceProvisioner implementation="com.google.gct.directaccess.provisioner.FirebaseDeviceProvisionerFactory"/>
+    <deviceProvisioner implementation="com.google.gct.directaccess.provisioner.DirectAccessDeviceProvisionerFactory"/>
   </extensions>
 </idea-plugin>
diff --git a/directaccess/src/com/google/gct/directaccess/DirectAccessService.kt b/directaccess/src/com/google/gct/directaccess/DirectAccessService.kt
index 770c2f7..7f02d61 100644
--- a/directaccess/src/com/google/gct/directaccess/DirectAccessService.kt
+++ b/directaccess/src/com/google/gct/directaccess/DirectAccessService.kt
@@ -15,7 +15,6 @@
  */
 package com.google.gct.directaccess
 
-import com.android.tools.adbbridge.Reservation
 import com.android.tools.idea.adblib.AdbLibService
 import com.android.tools.idea.concurrency.AndroidCoroutineScope
 import com.android.tools.idea.flags.StudioFlags
@@ -25,7 +24,6 @@ import com.google.gct.login.GoogleLogin
 import com.google.services.firebase.directaccess.client.DirectAccessConnection
 import com.google.services.firebase.directaccess.client.DirectAccessConnectionManager
 import com.google.services.firebase.directaccess.client.DirectAccessReservationManager
-import com.google.services.firebase.directaccess.client.isClosed
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.project.Project
@@ -71,33 +69,12 @@ class DirectAccessService(val project: Project) : Disposable {
     }
 
   /**
-   * Returns a [DirectAccessConnection] connecting to the remote device with [codename] and [api].
-   *
-   * The remote device is managed by a newly created [Reservation] from [reservationManager].
-   * However, if a [Reservation] with the same device information is created from another studio
-   * instance before calling this method, the existing [Reservation] will be reused by
-   * [reservationManager] and assigned to the returned [DirectAccessConnection].
+   * Returns a [DirectAccessConnection] connecting to a device with [reservationName] and [scope].
    */
-  fun reserveConnection(
-    codename: String,
-    api: String,
+  fun connectToReservation(
+    reservationName: String,
     scope: CoroutineScope
-  ): DirectAccessConnection? {
-    // These should be present if we are logged in.
-    val reservationManager = reservationManager ?: return null
-    val connectionManager = connectionManager ?: return null
-
-    val reservation =
-      reservationManager.listReservations().firstOrNull { reservation ->
-        !reservation.sessionState.isClosed() &&
-          reservation.androidDeviceList.androidDevicesList.any {
-            it.androidModelId == codename && it.androidVersionId == api
-          }
-      }
-        ?: reservationManager.createReservation(codename, api)
-
-    return connectionManager.connect(reservation, scope)
-  }
+  ): DirectAccessConnection? = connectionManager?.connect(reservationName, scope)
 
   override fun dispose() {}
 }
diff --git a/directaccess/src/com/google/gct/directaccess/provisioner/DirectAccessDeviceHandle.kt b/directaccess/src/com/google/gct/directaccess/provisioner/DirectAccessDeviceHandle.kt
new file mode 100644
index 0000000..706f2e7
--- /dev/null
+++ b/directaccess/src/com/google/gct/directaccess/provisioner/DirectAccessDeviceHandle.kt
@@ -0,0 +1,242 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * 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
+ *
+ *      http://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.
+ */
+package com.google.gct.directaccess.provisioner
+
+import com.android.adblib.ConnectedDevice
+import com.android.adblib.deviceProperties
+import com.android.sdklib.AndroidVersion
+import com.android.sdklib.deviceprovisioner.ActivationAction
+import com.android.sdklib.deviceprovisioner.ActivationParams
+import com.android.sdklib.deviceprovisioner.DeactivationAction
+import com.android.sdklib.deviceprovisioner.DeviceActionException
+import com.android.sdklib.deviceprovisioner.DeviceHandle
+import com.android.sdklib.deviceprovisioner.DeviceProperties
+import com.android.sdklib.deviceprovisioner.DeviceState
+import com.android.sdklib.deviceprovisioner.DeviceTemplate
+import com.android.sdklib.deviceprovisioner.ReservationAction
+import com.android.sdklib.deviceprovisioner.ReservationState
+import com.android.sdklib.deviceprovisioner.asMap
+import com.android.sdklib.deviceprovisioner.invokeOnDisconnection
+import com.android.tools.adbbridge.Reservation
+import com.android.tools.idea.devicemanager.DeviceType
+import com.android.tools.idea.run.DeviceHeadsUpListener
+import com.google.gct.directaccess.DirectAccessService
+import com.google.services.firebase.directaccess.client.DirectAccessConnection
+import com.google.services.firebase.directaccess.client.DirectAccessReservationManager
+import com.google.services.firebase.directaccess.client.isClosed
+import com.intellij.openapi.components.service
+import com.intellij.openapi.project.Project
+import java.time.Duration
+import java.time.Instant
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.NonCancellable
+import kotlinx.coroutines.TimeoutCancellationException
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.collect
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.flow.takeWhile
+import kotlinx.coroutines.flow.update
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
+import kotlinx.coroutines.withTimeout
+
+private val EXTENSION_TIMEOUT = Duration.ofSeconds(10)
+
+class DirectAccessDeviceHandle(
+  private val project: Project,
+  override val scope: CoroutineScope,
+  override val sourceTemplate: DeviceTemplate,
+  initialState: DeviceState,
+  reservationName: String,
+) : DeviceHandle {
+
+  private val reservationManager: DirectAccessReservationManager =
+    project.service<DirectAccessService>().reservationManager
+      ?: throw RuntimeException("Not logged in.")
+
+  val connection: DirectAccessConnection =
+    project.service<DirectAccessService>().connectToReservation(reservationName, scope)
+      ?: throw RuntimeException("Not logged in.")
+
+  override val stateFlow = MutableStateFlow(initialState)
+
+  init {
+    scope.launch {
+      // Map Reservation to its device provisioner format.
+      reservationManager
+        .fetchReservationFlow(reservationName)
+        .map {
+          val reservationState =
+            when (it.sessionState) {
+              Reservation.SessionState.REQUESTED,
+              Reservation.SessionState.PENDING -> ReservationState.PENDING
+              Reservation.SessionState.ACTIVE -> ReservationState.ACTIVE
+              Reservation.SessionState.FINISHED -> ReservationState.COMPLETE
+              else -> ReservationState.ERROR
+            }
+          com.android.sdklib.deviceprovisioner.Reservation(
+            reservationState,
+            "None",
+            Instant.ofEpochSecond(it.createTime.seconds),
+            Instant.ofEpochSecond(it.expireTime.seconds)
+          )
+        }
+        .collect { reservation ->
+          stateFlow.update { state ->
+            when (state) {
+              is DeviceState.Connected -> state.copy(reservation = reservation)
+              is DeviceState.Disconnected -> state.copy(reservation = reservation)
+              else -> state
+            }
+          }
+        }
+    }
+  }
+
+  override val activationAction =
+    object : ActivationAction {
+      /** Starts connection to the remote device. */
+      override suspend fun activate(params: ActivationParams) {
+        withContext(scope.coroutineContext) {
+          stateFlow.update { Activating(it.properties, it.reservation) }
+          connection.connect()
+          // Add disambiguator field that adds the port on which the device is connected to denote
+          // this is a firebase device.
+          // TODO(b/260153322): Remove once device manager moves to device provisioner framework
+          stateFlow.update {
+            Activating(
+              DirectAccessDeviceProperties.build {
+                manufacturer = it.properties.manufacturer
+                androidVersion = it.properties.androidVersion
+                model = it.properties.model
+                disambiguator = "${connection.port}"
+              },
+              it.reservation
+            )
+          }
+        }
+      }
+
+      override val label: String = "Connect"
+      override val isEnabled: StateFlow<Boolean> =
+        connection.state
+          .map { it.connection == DirectAccessConnection.ConnectionState.DISCONNECTED }
+          .stateIn(scope, SharingStarted.Eagerly, true)
+    }
+
+  override val deactivationAction =
+    object : DeactivationAction {
+      override suspend fun deactivate() =
+        withContext(scope.coroutineContext + NonCancellable) { connection.endReservation() }
+
+      override val label: String
+        get() = "Disconnect"
+      override val isEnabled: StateFlow<Boolean> =
+        connection.state
+          .map { !it.reservation.sessionState.isClosed() }
+          .stateIn(scope, SharingStarted.Eagerly, true)
+    }
+
+  override val reservationAction: ReservationAction =
+    object : ReservationAction {
+      override suspend fun reserve(duration: Duration): Instant {
+        val reservation =
+          state.reservation ?: throw DeviceActionException("Reservation not available.")
+        val endTime =
+          reservation.endTime ?: throw DeviceActionException("Reservation end time not available.")
+        connection.extendReservation(duration)
+        // Wait until reservation updates.
+        try {
+          withTimeout(EXTENSION_TIMEOUT.toMillis()) {
+            stateFlow
+              .takeWhile { it.reservation?.endTime?.toEpochMilli() == endTime.toEpochMilli() }
+              .collect()
+          }
+        } catch (e: TimeoutCancellationException) {
+          throw DeviceActionException(
+            "Reservation not extended within ${EXTENSION_TIMEOUT.seconds} seconds"
+          )
+        }
+        return state.reservation?.endTime
+          ?: throw DeviceActionException("Extended reservation end time not available.")
+      }
+
+      override val label: String = "Reserve"
+
+      /** [ReservationAction] is enabled through the lifecycle of the device handle. */
+      override val isEnabled: StateFlow<Boolean> = MutableStateFlow(true)
+    }
+
+  /** Returns true and changes state to [Connected] if [port] matches the [connection] of handle. */
+  suspend fun claim(port: Int, device: ConnectedDevice): Boolean {
+    if (connection.port != port) {
+      return false
+    }
+    // Show the device tab in running devices window.
+    project.messageBus
+      .syncPublisher(DeviceHeadsUpListener.TOPIC)
+      .userInvolvementRequired(device.deviceInfoFlow.value.serialNumber, project)
+    val properties = device.deviceProperties().all().asMap()
+    val deviceProperties =
+      DirectAccessDeviceProperties.build {
+        readCommonProperties(properties)
+        // TODO(b/260153322): Remove once device manager moves to device provisioner framework
+        disambiguator = "${connection.port}"
+      }
+    stateFlow.update { DeviceState.Connected(deviceProperties, device, it.reservation) }
+    device.invokeOnDisconnection {
+      stateFlow.update {
+        DeviceState.Disconnected(deviceProperties, false, it.status, it.reservation)
+      }
+    }
+    return true
+  }
+
+  class Activating(
+    override val properties: DeviceProperties,
+    reservation: com.android.sdklib.deviceprovisioner.Reservation?
+  ) : DeviceState.Disconnected(properties, isTransitioning = true, "Connecting", reservation)
+}
+
+class DirectAccessDeviceProperties(base: DeviceProperties) : DeviceProperties by base {
+  class Builder : DeviceProperties.Builder()
+  companion object {
+    fun build(block: Builder.() -> Unit) =
+      Builder().apply(block).run { DirectAccessDeviceProperties(buildBase()) }
+  }
+}
+
+fun DeviceInfo.toDeviceProperties(): DirectAccessDeviceProperties {
+  val info = this
+  return DirectAccessDeviceProperties.build {
+    manufacturer = info.manufacturer
+    model = info.name
+    androidVersion = AndroidVersion(info.api)
+    deviceType =
+      when (info.type) {
+        DeviceType.PHONE -> com.android.sdklib.deviceprovisioner.DeviceType.HANDHELD
+        DeviceType.TV -> com.android.sdklib.deviceprovisioner.DeviceType.TV
+        DeviceType.WEAR_OS -> com.android.sdklib.deviceprovisioner.DeviceType.WEAR
+        DeviceType.AUTOMOTIVE -> com.android.sdklib.deviceprovisioner.DeviceType.AUTOMOTIVE
+      }
+  }
+}
+
+fun ReservationState.isClosed() =
+  this == ReservationState.ERROR || this == ReservationState.COMPLETE
diff --git a/directaccess/src/com/google/gct/directaccess/provisioner/FirebaseDeviceProvisionerFactory.kt b/directaccess/src/com/google/gct/directaccess/provisioner/DirectAccessDeviceProvisionerFactory.kt
index af88ab7..1e0e38b 100644
--- a/directaccess/src/com/google/gct/directaccess/provisioner/FirebaseDeviceProvisionerFactory.kt
+++ b/directaccess/src/com/google/gct/directaccess/provisioner/DirectAccessDeviceProvisionerFactory.kt
@@ -21,10 +21,10 @@ import com.android.tools.idea.flags.StudioFlags
 import com.intellij.openapi.project.Project
 import kotlinx.coroutines.CoroutineScope
 
-class FirebaseDeviceProvisionerFactory : DeviceProvisionerFactory {
+class DirectAccessDeviceProvisionerFactory : DeviceProvisionerFactory {
   override val isEnabled: Boolean
     get() = StudioFlags.DIRECT_ACCESS.get()
 
-  override fun create(scope: CoroutineScope, project: Project): DeviceProvisionerPlugin =
-    FirebaseDeviceProvisioner(scope, project)
+  override fun create(coroutineScope: CoroutineScope, project: Project): DeviceProvisionerPlugin =
+    DirectAccessDeviceProvisionerPlugin(coroutineScope, project)
 }
diff --git a/directaccess/src/com/google/gct/directaccess/provisioner/DirectAccessDeviceProvisionerPlugin.kt b/directaccess/src/com/google/gct/directaccess/provisioner/DirectAccessDeviceProvisionerPlugin.kt
new file mode 100644
index 0000000..84f8d48
--- /dev/null
+++ b/directaccess/src/com/google/gct/directaccess/provisioner/DirectAccessDeviceProvisionerPlugin.kt
@@ -0,0 +1,158 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * 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
+ *
+ *      http://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.
+ */
+package com.google.gct.directaccess.provisioner
+
+import com.android.adblib.ConnectedDevice
+import com.android.sdklib.deviceprovisioner.DeviceHandle
+import com.android.sdklib.deviceprovisioner.DeviceProvisionerPlugin
+import com.android.sdklib.deviceprovisioner.DeviceTemplate
+import com.android.tools.idea.concurrency.createChildScope
+import com.android.tools.idea.flags.StudioFlags
+import com.google.gct.directaccess.DirectAccessService
+import com.google.gct.login.LoginState
+import com.google.services.firebase.directaccess.client.isClosed
+import com.intellij.openapi.components.service
+import com.intellij.openapi.diagnostic.Logger
+import com.intellij.openapi.project.Project
+import java.util.concurrent.TimeUnit
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.launch
+
+private val defaultDeviceInfoProvider = {
+  CatalogClient.getAvailableDevices("https://${StudioFlags.DIRECT_ACCESS_ENDPOINT.get()}/")
+}
+
+/**
+ * Provides direct access to physical devices run by Firebase. Supports configuring direct access
+ * device templates and activating / deactivating them.
+ */
+class DirectAccessDeviceProvisionerPlugin(
+  private val scope: CoroutineScope,
+  private val project: Project,
+  private val deviceInfoProvider: () -> List<DeviceInfo> = defaultDeviceInfoProvider
+) : DeviceProvisionerPlugin {
+  private val logger = Logger.getInstance(DirectAccessDeviceProvisionerPlugin::class.java)
+
+  // TODO: find a proper priority
+  override val priority: Int = 120
+
+  private val _devices = MutableStateFlow(emptyList<DeviceHandle>())
+  override val devices: StateFlow<List<DeviceHandle>> = _devices
+  private val _templates = MutableStateFlow(emptyList<DeviceTemplate>())
+  override val templates: StateFlow<List<DeviceTemplate>> = _templates
+
+  init {
+    // This scope will not be cancelled on login changes. Only the inner child scope will be
+    // cancelled.
+    scope.launch {
+      var childScope: CoroutineScope? = null
+      LoginState.loggedIn.distinctUntilChanged().collect { isLoggedIn ->
+        // This cancellation will cause all child scopes created from this childScope to be
+        // cancelled.
+        // This includes cancellation of scopes in template, handle, connection.
+        childScope?.cancel()
+        childScope = scope.createChildScope(isSupervisor = true)
+        if (isLoggedIn) {
+          childScope?.launch { periodicUpdateReservation() }
+          childScope?.launch {
+            // Fetch reservations with new templates.
+            templates.collect { updateReservations(project, templates) }
+          }
+          childScope?.launch { periodicUpdateTemplates(this) }
+        } else {
+          _templates.value = listOf()
+        }
+      }
+    }
+  }
+
+  // Update templates every 5 minutes.
+  private suspend fun periodicUpdateTemplates(parentScope: CoroutineScope) {
+    while (true) {
+      val oldTemplates =
+        templates.value
+          .groupBy { (it as DirectAccessDeviceTemplate).deviceInfo }
+          .mapValues { it.value[0] }
+      try {
+        deviceInfoProvider()
+          .map { info ->
+            // Create a child scope for every template to isolate them in terms of scope.
+            // This helps avoid any issues with a given template from propagating to other
+            // templates.
+            oldTemplates[info]
+              ?: DirectAccessDeviceTemplate(
+                project,
+                info,
+                _devices,
+                parentScope.createChildScope(isSupervisor = true)
+              )
+          }
+          .let { result -> _templates.value = result }
+      } catch (ignore: NotLoggedInException) {
+        // do nothing
+      } catch (e: Exception) {
+        logger.warn(e)
+      }
+      delay(TimeUnit.MINUTES.toMillis(5))
+    }
+  }
+
+  // Fetch reservations periodically in case a Reservation is created elsewhere.
+  private suspend fun periodicUpdateReservation() {
+    while (true) {
+      updateReservations(project, templates)
+      delay(TimeUnit.MINUTES.toMillis(1))
+    }
+  }
+
+  override suspend fun claim(device: ConnectedDevice): DeviceHandle? {
+    val sn = device.deviceInfoFlow.value.serialNumber
+    if (sn.matches(Regex("^localhost:\\d+$"))) {
+      val port = sn.substringAfter(':').toIntOrNull() ?: return null
+      return devices.value.filterIsInstance<DirectAccessDeviceHandle>().firstOrNull {
+        it.claim(port, device)
+      }
+    }
+    return null
+  }
+}
+
+suspend fun updateReservations(project: Project, templates: StateFlow<List<DeviceTemplate>>) {
+  val templateMap =
+    templates.value.filterIsInstance<DirectAccessDeviceTemplate>().groupBy { template ->
+      template.deviceInfo.let { "${it.codename} ${it.api}" }
+    }
+  project
+    .service<DirectAccessService>()
+    .reservationManager
+    ?.listReservations()
+    ?.filter { reservation ->
+      !reservation.sessionState.isClosed() &&
+        reservation.androidDeviceList.androidDevicesList.isNotEmpty()
+    }
+    ?.forEach { reservation ->
+      val key =
+        reservation.androidDeviceList.androidDevicesList[0].let {
+          "${it.androidModelId} ${it.androidVersionId}"
+        }
+      templateMap[key]?.firstOrNull()?.activationAction?.activate()
+    }
+}
diff --git a/directaccess/src/com/google/gct/directaccess/provisioner/DirectAccessDeviceTemplate.kt b/directaccess/src/com/google/gct/directaccess/provisioner/DirectAccessDeviceTemplate.kt
new file mode 100644
index 0000000..9ec489e
--- /dev/null
+++ b/directaccess/src/com/google/gct/directaccess/provisioner/DirectAccessDeviceTemplate.kt
@@ -0,0 +1,124 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * 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
+ *
+ *      http://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.
+ */
+package com.google.gct.directaccess.provisioner
+
+import com.android.sdklib.deviceprovisioner.DeviceActionDisabledException
+import com.android.sdklib.deviceprovisioner.DeviceActionException
+import com.android.sdklib.deviceprovisioner.DeviceHandle
+import com.android.sdklib.deviceprovisioner.DeviceState
+import com.android.sdklib.deviceprovisioner.DeviceTemplate
+import com.android.sdklib.deviceprovisioner.TemplateActivationAction
+import com.android.tools.idea.concurrency.createChildScope
+import com.google.gct.directaccess.DirectAccessService
+import com.google.services.firebase.directaccess.client.findOrCreateReservation
+import com.intellij.openapi.components.service
+import com.intellij.openapi.project.Project
+import java.time.Duration
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.update
+import kotlinx.coroutines.launch
+
+class DirectAccessDeviceTemplate(
+  private val project: Project,
+  val deviceInfo: DeviceInfo,
+  private val devices: MutableStateFlow<List<DeviceHandle>>,
+  private val scope: CoroutineScope
+) : DeviceTemplate {
+  override val properties = deviceInfo.toDeviceProperties()
+
+  private val isActivationEnabled = MutableStateFlow(true)
+
+  /**
+   * Last device handle activated by the template.
+   *
+   * TODO (b/246171065): resolve potential race condition to support activating multiple devices
+   */
+  var activeDevice: DirectAccessDeviceHandle? = null
+    private set(device) {
+      field = device
+      if (device != null) {
+        devices.update { list -> list + device }
+        device.scope.launch {
+          device.stateFlow.collect {
+            if (it.reservation?.state?.isClosed() == true) {
+              field = null
+              isActivationEnabled.value = true
+              devices.update { list -> list - device }
+            }
+          }
+        }
+      }
+    }
+
+  override val activationAction: TemplateActivationAction =
+    object : TemplateActivationAction {
+      // TODO: Pass duration through to the DirectAccessConnectionManager.
+      override val durationUsed = false
+
+      /**
+       * Creates a [DirectAccessDeviceHandle] with [Disconnected] state.
+       *
+       * This method first finds or create a [Reservation] that matches its [deviceInfo]. Then a
+       * [DirectAccessDeviceHandle] is created with a [DirectAccessConnection] to the [Reservation].
+       * Connection is not started until the activationAction of device handle get called. The
+       * device handle is added to the devices flow of the provisioner and will be removed after
+       * [Reservation] closed. This method is disabled when a device handle is activating or
+       * activated. At most one device is available for each template.
+       *
+       * TODO (b/246171065): activating multiple devices.
+       */
+      override suspend fun activate(duration: Duration?): DeviceHandle {
+        // Disable further activate actions to avoid multiple devices.
+        if (!isActivationEnabled.compareAndSet(expect = true, update = false)) {
+          throw DeviceActionDisabledException(this)
+        }
+
+        val reservationManager =
+          project.service<DirectAccessService>().reservationManager
+            ?: throw DeviceActionException("Unable to access ReservationManager.")
+
+        val reservationName =
+          try {
+            reservationManager
+              .findOrCreateReservation(deviceInfo.codename, deviceInfo.api.toString())
+              .name
+          } catch (e: Exception) {
+            // Pass the underlying gRPC exception as a cause
+            // TODO: Perhaps extract more detail if we can get it.
+            throw DeviceActionException("Unable to reserve device.", e)
+          }
+
+        val deviceProperties = deviceInfo.toDeviceProperties()
+        val deviceScope = scope.createChildScope(isSupervisor = true)
+        // Notify provisioner plugin of the new device.
+        return DirectAccessDeviceHandle(
+            project,
+            deviceScope,
+            this@DirectAccessDeviceTemplate,
+            DeviceState.Disconnected(deviceProperties),
+            reservationName
+          )
+          .also { activeDevice = it }
+      }
+
+      override val label: String = "Acquire"
+      override val isEnabled: StateFlow<Boolean> = isActivationEnabled
+    }
+
+  override val editAction = null
+}
diff --git a/directaccess/src/com/google/gct/directaccess/provisioner/FirebaseDeviceProvisioner.kt b/directaccess/src/com/google/gct/directaccess/provisioner/FirebaseDeviceProvisioner.kt
deleted file mode 100644
index 7aad46a..0000000
--- a/directaccess/src/com/google/gct/directaccess/provisioner/FirebaseDeviceProvisioner.kt
+++ /dev/null
@@ -1,359 +0,0 @@
-/*
- * Copyright (C) 2022 The Android Open Source Project
- *
- * 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
- *
- *      http://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.
- */
-package com.google.gct.directaccess.provisioner
-
-import com.android.adblib.ConnectedDevice
-import com.android.adblib.deviceProperties
-import com.android.sdklib.AndroidVersion
-import com.android.sdklib.deviceprovisioner.ActivationAction
-import com.android.sdklib.deviceprovisioner.ActivationParams
-import com.android.sdklib.deviceprovisioner.DeactivationAction
-import com.android.sdklib.deviceprovisioner.DeviceActionDisabledException
-import com.android.sdklib.deviceprovisioner.DeviceActionException
-import com.android.sdklib.deviceprovisioner.DeviceHandle
-import com.android.sdklib.deviceprovisioner.DeviceProperties
-import com.android.sdklib.deviceprovisioner.DeviceProvisionerPlugin
-import com.android.sdklib.deviceprovisioner.DeviceState
-import com.android.sdklib.deviceprovisioner.DeviceState.Connected
-import com.android.sdklib.deviceprovisioner.DeviceState.Disconnected
-import com.android.sdklib.deviceprovisioner.DeviceTemplate
-import com.android.sdklib.deviceprovisioner.TemplateActivationAction
-import com.android.sdklib.deviceprovisioner.asMap
-import com.android.sdklib.deviceprovisioner.invokeOnDisconnection
-import com.android.tools.adbbridge.Reservation
-import com.android.tools.idea.concurrency.createChildScope
-import com.android.tools.idea.flags.StudioFlags
-import com.android.tools.idea.run.DeviceHeadsUpListener
-import com.google.gct.directaccess.DirectAccessService
-import com.google.gct.login.LoginState
-import com.google.services.firebase.directaccess.client.DirectAccessConnection
-import com.google.services.firebase.directaccess.client.isClosed
-import com.intellij.openapi.components.service
-import com.intellij.openapi.diagnostic.Logger
-import com.intellij.openapi.project.Project
-import java.time.Duration
-import java.util.concurrent.TimeUnit
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.NonCancellable
-import kotlinx.coroutines.cancel
-import kotlinx.coroutines.delay
-import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.SharingStarted
-import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.distinctUntilChanged
-import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.flow.stateIn
-import kotlinx.coroutines.flow.update
-import kotlinx.coroutines.launch
-import kotlinx.coroutines.withContext
-
-private val defaultDeviceInfoProvider = {
-  CatalogClient.getAvailableDevices("https://${StudioFlags.DIRECT_ACCESS_ENDPOINT.get()}/")
-}
-
-/**
- * Provides access to physical devices run by Firebase. Supports configuring Firebase device
- * templates and activating / deactivating them.
- */
-class FirebaseDeviceProvisioner(
-  private val scope: CoroutineScope,
-  private val project: Project,
-  private val deviceInfoProvider: () -> List<DeviceInfo> = defaultDeviceInfoProvider
-) : DeviceProvisionerPlugin {
-  private val logger = Logger.getInstance(FirebaseDeviceProvisioner::class.java)
-
-  // TODO: find a proper priority
-  override val priority: Int = 120
-
-  private val _devices = MutableStateFlow(emptyList<DeviceHandle>())
-  override val devices: StateFlow<List<DeviceHandle>> = _devices
-  private val _templates = MutableStateFlow(emptyList<DeviceTemplate>())
-  override val templates: StateFlow<List<DeviceTemplate>> = _templates
-
-  init {
-    // This scope will not be cancelled on login changes. Only the inner child scope will be
-    // cancelled.
-    scope.launch {
-      var childScope: CoroutineScope? = null
-      LoginState.loggedIn.distinctUntilChanged().collect { isLoggedIn ->
-        // This cancellation will cause all child scopes created from this childScope to be
-        // cancelled.
-        // This includes cancellation of scopes in template, handle, connection.
-        childScope?.cancel()
-        childScope = scope.createChildScope(isSupervisor = true)
-        if (isLoggedIn) {
-          childScope?.launch { periodicUpdateReservation() }
-          childScope?.launch {
-            // Fetch reservations with new templates.
-            templates.collect { updateReservations(project, templates) }
-          }
-          childScope?.launch { periodicUpdateTemplates(this) }
-        }
-      }
-    }
-  }
-
-  // Update templates every 5 minutes.
-  private suspend fun periodicUpdateTemplates(parentScope: CoroutineScope) {
-    while (true) {
-      val oldTemplates =
-        templates.value
-          .groupBy { (it as FirebaseDeviceTemplate).deviceInfo }
-          .mapValues { it.value[0] }
-      try {
-        deviceInfoProvider()
-          .map { info ->
-            // Create a child scope for every template to isolate them in terms of scope.
-            // This helps avoid any issues with a given template from propagating to other
-            // templates.
-            oldTemplates[info]
-              ?: FirebaseDeviceTemplate(
-                project,
-                info,
-                _devices,
-                parentScope.createChildScope(isSupervisor = true)
-              )
-          }
-          .let { result -> _templates.value = result }
-      } catch (ignore: NotLoggedInException) {
-        // do nothing
-      } catch (e: Exception) {
-        logger.warn(e)
-      }
-      delay(TimeUnit.MINUTES.toMillis(5))
-    }
-  }
-
-  // Fetch reservations periodically in case a Reservation is created elsewhere.
-  private suspend fun periodicUpdateReservation() {
-    while (true) {
-      updateReservations(project, templates)
-      delay(TimeUnit.MINUTES.toMillis(1))
-    }
-  }
-
-  override suspend fun claim(device: ConnectedDevice): DeviceHandle? {
-    val sn = device.deviceInfoFlow.value.serialNumber
-    if (sn.matches(Regex("^localhost:\\d+$"))) {
-      val port = sn.substringAfter(':').toIntOrNull() ?: return null
-      return devices.value.filterIsInstance<DirectAccessDeviceHandle>().firstOrNull {
-        it.claim(port, device)
-      }
-    }
-    return null
-  }
-}
-
-suspend fun updateReservations(project: Project, templates: StateFlow<List<DeviceTemplate>>) {
-  val templateMap =
-    templates.value.filterIsInstance<FirebaseDeviceTemplate>().groupBy { template ->
-      template.deviceInfo.let { "${it.codename} ${it.api}" }
-    }
-  project
-    .service<DirectAccessService>()
-    .reservationManager
-    ?.listReservations()
-    ?.filter { reservation ->
-      !reservation.sessionState.isClosed() &&
-        reservation.androidDeviceList.androidDevicesList.isNotEmpty()
-    }
-    ?.forEach { reservation ->
-      val key =
-        reservation.androidDeviceList.androidDevicesList[0].let {
-          "${it.androidModelId} ${it.androidVersionId}"
-        }
-      templateMap[key]?.firstOrNull()?.activationAction?.activate()
-    }
-}
-
-class FirebaseDeviceTemplate(
-  private val project: Project,
-  val deviceInfo: DeviceInfo,
-  devices: MutableStateFlow<List<DeviceHandle>>,
-  private val scope: CoroutineScope
-) : DeviceTemplate {
-  override val displayName: String = "${deviceInfo.manufacturer} ${deviceInfo.name}"
-
-  /**
-   * Last device handle activated by the template.
-   *
-   * TODO (b/246171065): resolve potential race condition to support activating multiple devices
-   */
-  var activeDevice: DirectAccessDeviceHandle? = null
-    private set
-
-  override val activationAction: TemplateActivationAction =
-    object : TemplateActivationAction {
-      private val _isEnabled = MutableStateFlow(true)
-
-      // TODO: Pass duration through to the DirectAccessConnectionManager.
-      override val durationUsed = false
-
-      /**
-       * Creates a [DirectAccessDeviceHandle] with [Disconnected] state.
-       *
-       * This method first finds or create a [Reservation] that matches its [deviceInfo]. Then a
-       * [DirectAccessDeviceHandle] is created with a [DirectAccessConnection] to the [Reservation].
-       * Connection is not started until the activationAction of device handle get called. The
-       * device handle is added to the devices flow of the provisioner and will be removed after
-       * [Reservation] closed. This method is disabled when a device handle is activating or
-       * activated. At most one device is available for each template.
-       *
-       * TODO (b/246171065): activating multiple devices.
-       */
-      override suspend fun activate(duration: Duration?): DeviceHandle {
-        // Disable further activate actions to avoid multiple devices.
-        if (!_isEnabled.compareAndSet(expect = true, update = false)) {
-          throw DeviceActionDisabledException(this)
-        }
-
-        val connection =
-          try {
-            project
-              .service<DirectAccessService>()
-              .reserveConnection(deviceInfo.codename, deviceInfo.api.toString(), scope)
-              ?: throw DeviceActionException("Unable to reserve device.")
-          } catch (e: Exception) {
-            // Pass the underlying gRPC exception as a cause
-            // TODO: Perhaps extract more detail if we can get it.
-            throw DeviceActionException("Unable to reserve device.", e)
-          }
-
-        val deviceProperties =
-          DirectAccessDeviceProperties.build {
-            manufacturer = deviceInfo.manufacturer
-            androidVersion = AndroidVersion(deviceInfo.api)
-            model = deviceInfo.name
-          }
-        val deviceScope = scope.createChildScope(isSupervisor = true)
-
-        return DirectAccessDeviceHandle(
-            project,
-            deviceScope,
-            Disconnected(deviceProperties),
-            connection
-          )
-          .also { device ->
-            activeDevice = device
-            // Notify provisioner plugin of the new device.
-            devices.update { list -> list + device }
-            deviceScope.launch {
-              connection.state.collect {
-                if (it.reservation.sessionState.isClosed()) {
-                  activeDevice = null
-                  _isEnabled.value = true
-                  devices.update { list -> list - device }
-                }
-              }
-            }
-          }
-      }
-
-      override val label: String = "Acquire"
-      override val isEnabled: StateFlow<Boolean> = _isEnabled
-    }
-
-  override val editAction = null
-}
-
-class DirectAccessDeviceHandle(
-  private val project: Project,
-  override val scope: CoroutineScope,
-  state: DeviceState,
-  val connection: DirectAccessConnection
-) : DeviceHandle {
-
-  override val stateFlow = MutableStateFlow(state)
-
-  override val activationAction =
-    object : ActivationAction {
-      /** Starts connection to the remote device. */
-      override suspend fun activate(params: ActivationParams) {
-        withContext(scope.coroutineContext) {
-          stateFlow.update { Activating(it.properties) }
-          connection.connect()
-          // Add disambiguator field that adds the port on which the device is connected to denote
-          // this is a firebase device.
-          // TODO(b/260153322): Remove once device manager moves to device provisioner framework
-          stateFlow.update {
-            Activating(
-              DirectAccessDeviceProperties.build {
-                manufacturer = it.properties.manufacturer
-                androidVersion = it.properties.androidVersion
-                model = it.properties.model
-                disambiguator = "${connection.port}"
-              }
-            )
-          }
-        }
-      }
-
-      override val label: String = "Connect"
-      override val isEnabled: StateFlow<Boolean> =
-        connection.state
-          .map { it.connection == DirectAccessConnection.ConnectionState.DISCONNECTED }
-          .stateIn(scope, SharingStarted.Eagerly, true)
-    }
-
-  override val deactivationAction =
-    object : DeactivationAction {
-      override suspend fun deactivate() {
-        withContext(scope.coroutineContext + NonCancellable) {
-          connection.endReservation()
-          stateFlow.value = Disconnected(stateFlow.value.properties)
-        }
-      }
-
-      override val label: String
-        get() = "Disconnect"
-      override val isEnabled: StateFlow<Boolean> =
-        connection.state
-          .map { !it.reservation.sessionState.isClosed() }
-          .stateIn(scope, SharingStarted.Eagerly, true)
-    }
-
-  /** Returns true and changes state to [Connected] if [port] matches the [connection] of handle. */
-  suspend fun claim(port: Int, device: ConnectedDevice): Boolean {
-    if (connection.port != port) {
-      return false
-    }
-    // Show the device tab in running devices window.
-    project.messageBus
-      .syncPublisher(DeviceHeadsUpListener.TOPIC)
-      .userInvolvementRequired(device.deviceInfoFlow.value.serialNumber, project)
-    val properties = device.deviceProperties().all().asMap()
-    val deviceProperties =
-      DirectAccessDeviceProperties.build {
-        readCommonProperties(properties)
-        // TODO(b/260153322): Remove once device manager moves to device provisioner framework
-        disambiguator = "${connection.port}"
-      }
-    stateFlow.value = Connected(deviceProperties, device)
-    device.invokeOnDisconnection { stateFlow.value = Disconnected(deviceProperties) }
-    return true
-  }
-
-  class Activating(override val properties: DeviceProperties) :
-    Disconnected(properties, isTransitioning = true, "Connecting")
-}
-
-class DirectAccessDeviceProperties(base: DeviceProperties) : DeviceProperties by base {
-  class Builder : DeviceProperties.Builder()
-  companion object {
-    fun build(block: Builder.() -> Unit) =
-      Builder().apply(block).run { DirectAccessDeviceProperties(buildBase()) }
-  }
-}
diff --git a/directaccess/src/com/google/gct/directaccess/ui/FirebaseDeviceTable.kt b/directaccess/src/com/google/gct/directaccess/ui/FirebaseDeviceTable.kt
index 9fc5f7a..4b057d8 100644
--- a/directaccess/src/com/google/gct/directaccess/ui/FirebaseDeviceTable.kt
+++ b/directaccess/src/com/google/gct/directaccess/ui/FirebaseDeviceTable.kt
@@ -75,7 +75,7 @@ class FirebaseDeviceTable(
         Comparator.comparing { item: FirebaseItem ->
           when (item) {
             is FirebaseDeviceItem -> item.device.name
-            is FirebaseDeviceTemplateItem -> item.template.displayName
+            is FirebaseDeviceTemplateItem -> item.template.properties.title
             else -> ""
           }
         }
diff --git a/directaccess/src/com/google/gct/directaccess/ui/FirebaseItemManager.kt b/directaccess/src/com/google/gct/directaccess/ui/FirebaseItemManager.kt
index fe07333..b640cad 100644
--- a/directaccess/src/com/google/gct/directaccess/ui/FirebaseItemManager.kt
+++ b/directaccess/src/com/google/gct/directaccess/ui/FirebaseItemManager.kt
@@ -23,9 +23,9 @@ import com.android.tools.idea.deviceprovisioner.DeviceProvisionerService
 import com.android.tools.idea.flags.StudioFlags
 import com.google.gct.directaccess.FirebaseDevice
 import com.google.gct.directaccess.provisioner.DirectAccessDeviceHandle
-import com.google.gct.directaccess.provisioner.FirebaseDeviceTemplate
+import com.google.gct.directaccess.provisioner.DirectAccessDeviceTemplate
+import com.google.gct.directaccess.provisioner.isClosed
 import com.google.services.firebase.directaccess.client.DirectAccessConnection.ConnectionState
-import com.google.services.firebase.directaccess.client.isClosed
 import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.ui.MessageDialogBuilder
@@ -121,7 +121,7 @@ class FirebaseDeviceItem(
 
 class FirebaseDeviceTemplateItem(
   private val itemManager: FirebaseItemManager,
-  val template: FirebaseDeviceTemplate,
+  val template: DirectAccessDeviceTemplate,
   private val scope: CoroutineScope,
   private val uiDispatcher: CoroutineDispatcher,
   override val onUpdate: () -> Unit
@@ -164,8 +164,8 @@ class FirebaseDeviceTemplateItem(
         newDevice?.let { newDeviceHandle ->
           scope.launch {
             newDeviceHandle.connection.state
-              .combine(newDevice.stateFlow) { remoteState, deviceState ->
-                if (remoteState.reservation.sessionState.isClosed()) {
+              .combine(newDeviceHandle.stateFlow) { remoteState, deviceState ->
+                if (deviceState.reservation?.state?.isClosed() == true) {
                   deviceItem = null
                   coroutineContext.cancel()
                 } else {
@@ -207,7 +207,7 @@ class FirebaseItemManager(
   init {
     scope.launch {
       provisionerPlugin.templates
-        .map { it.filterIsInstance<FirebaseDeviceTemplate>() }
+        .map { it.filterIsInstance<DirectAccessDeviceTemplate>() }
         .distinctUntilChanged()
         .collect { newTemplates -> refreshTemplates(newTemplates) }
     }
@@ -219,7 +219,7 @@ class FirebaseItemManager(
     }
   }
 
-  private suspend fun refreshTemplates(newTemplates: List<FirebaseDeviceTemplate>) {
+  private suspend fun refreshTemplates(newTemplates: List<DirectAccessDeviceTemplate>) {
     withContext(uiDispatcher) {
       val existingMap = templateItems.associateBy { it.template.deviceInfo }
       templateItems =
diff --git a/directaccess/src/com/google/gct/directaccess/ui/FirebaseTemplateTableCellRenderer.kt b/directaccess/src/com/google/gct/directaccess/ui/FirebaseTemplateTableCellRenderer.kt
index f4bfb80..54b1c9c 100644
--- a/directaccess/src/com/google/gct/directaccess/ui/FirebaseTemplateTableCellRenderer.kt
+++ b/directaccess/src/com/google/gct/directaccess/ui/FirebaseTemplateTableCellRenderer.kt
@@ -16,7 +16,7 @@
 package com.google.gct.directaccess.ui
 
 import com.android.tools.idea.devicemanager.Tables
-import com.google.gct.directaccess.provisioner.FirebaseDeviceTemplate
+import com.google.gct.directaccess.provisioner.DirectAccessDeviceTemplate
 import com.intellij.ui.components.JBLabel
 import com.intellij.ui.components.JBPanel
 import com.intellij.ui.scale.JBUIScale
@@ -96,10 +96,10 @@ class FirebaseTemplateTableCellRenderer : TableCellRenderer {
     viewRowIndex: Int,
     viewColumnIndex: Int
   ): Component {
-    val template = value as FirebaseDeviceTemplate
+    val template = value as DirectAccessDeviceTemplate
     val foreground = Tables.getForeground(table, selected)
     nameLabel.foreground = foreground
-    nameLabel.text = template.displayName
+    nameLabel.text = template.properties.title
     stateLabel.foreground = foreground
     line2Label.font = UIUtil.getLabelFont(UIUtil.FontSize.SMALL)
     line2Label.foreground = foreground.brighter()
diff --git a/directaccess/testSrc/com/google/gct/directaccess/provisioner/FirebaseDeviceProvisionerTest.kt b/directaccess/testSrc/com/google/gct/directaccess/provisioner/DirectAccessDeviceProvisionerTest.kt
index 4c71596..017d958 100644
--- a/directaccess/testSrc/com/google/gct/directaccess/provisioner/FirebaseDeviceProvisionerTest.kt
+++ b/directaccess/testSrc/com/google/gct/directaccess/provisioner/DirectAccessDeviceProvisionerTest.kt
@@ -34,33 +34,36 @@ import com.google.common.util.concurrent.MoreExecutors
 import com.google.gct.directaccess.DirectAccessService
 import com.google.gct.directaccess.TestUtils.deviceInfoListProvider
 import com.google.gct.login.GoogleLogin
+import com.google.gct.login.LoginState
 import com.google.services.firebase.directaccess.client.DirectAccessReservationManager
 import com.google.services.firebase.directaccess.client.FakeDirectAccessConnection
 import com.google.services.firebase.directaccess.client.FakeDirectAccessGrpcService
 import com.google.services.firebase.directaccess.client.deviceAddress
 import com.studiogrpc.testutils.GrpcConnectionRule
+import java.time.Duration
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.asCoroutineDispatcher
 import kotlinx.coroutines.cancel
+import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.launch
 import org.junit.After
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
-import org.mockito.Mockito.anyString
+import org.mockito.Mockito.doAnswer
 import org.mockito.Mockito.doReturn
 
-class FirebaseDeviceProvisionerTest {
+class DirectAccessDeviceProvisionerTest {
 
   private val service = FakeDirectAccessGrpcService()
   @get:Rule val projectRule = AndroidProjectRule.inMemory()
   @get:Rule val grpcConnectionRule = GrpcConnectionRule(listOf(service))
 
   private val session = FakeAdbSession()
-  private val fakeConnection = FakeDirectAccessConnection()
-  private lateinit var plugin: FirebaseDeviceProvisioner
+  private lateinit var plugin: DirectAccessDeviceProvisionerPlugin
   private lateinit var provisioner: DeviceProvisioner
   private lateinit var directAccessReservationManager: DirectAccessReservationManager
+  private lateinit var fakeConnection: FakeDirectAccessConnection
   private lateinit var scope: CoroutineScope
   private lateinit var mockGoogleLogin: GoogleLogin
 
@@ -68,17 +71,29 @@ class FirebaseDeviceProvisionerTest {
   fun setUp() = runBlockingWithTimeout {
     mockGoogleLogin = projectRule.mockService(GoogleLogin::class.java)
     doReturn(true).whenever(mockGoogleLogin).isLoggedIn
+    (LoginState.loggedIn as MutableStateFlow<Boolean>).value = true
     scope = CoroutineScope(MoreExecutors.directExecutor().asCoroutineDispatcher())
     directAccessReservationManager =
       DirectAccessReservationManager("testProject", scope, grpcConnectionRule.channel) {
         "testToken"
       }
     val mockDirectAccessService = projectRule.mockProjectService(DirectAccessService::class.java)
-    doReturn(fakeConnection)
-      .whenever(mockDirectAccessService)
-      .reserveConnection(anyString(), anyString(), any())
     doReturn(directAccessReservationManager).whenever(mockDirectAccessService).reservationManager
-    plugin = FirebaseDeviceProvisioner(session.scope, projectRule.project, deviceInfoListProvider)
+    doAnswer {
+        val reservationName = it.arguments[0] as String
+        val deviceScope = it.arguments[1] as CoroutineScope
+        fakeConnection =
+          FakeDirectAccessConnection(directAccessReservationManager, reservationName, deviceScope)
+        fakeConnection
+      }
+      .whenever(mockDirectAccessService)
+      .connectToReservation(any(), any())
+    plugin =
+      DirectAccessDeviceProvisionerPlugin(
+        session.scope,
+        projectRule.project,
+        deviceInfoListProvider
+      )
     provisioner = DeviceProvisioner.create(session, listOf(plugin))
     yieldUntil { provisioner.templates.value.isNotEmpty() }
   }
@@ -96,9 +111,13 @@ class FirebaseDeviceProvisionerTest {
     yieldUntil { provisioner.templates.value.size == 3 }
 
     // Assert
-    assertThat(provisioner.templates.value[0].displayName).isEqualTo("Google Pixel 5")
-    assertThat(provisioner.templates.value[1].displayName).isEqualTo("Google Pixel 6")
-    assertThat(provisioner.templates.value[2].displayName).isEqualTo("Google Pixel 6 Pro")
+    assertThat(provisioner.templates.value[0].properties.title).isEqualTo("Google Pixel 5")
+    assertThat(provisioner.templates.value[1].properties.title).isEqualTo("Google Pixel 6")
+    assertThat(provisioner.templates.value[2].properties.title).isEqualTo("Google Pixel 6 Pro")
+
+    // Log out
+    (LoginState.loggedIn as MutableStateFlow<Boolean>).value = false
+    yieldUntil { provisioner.templates.value.isEmpty() }
   }
 
   @Test
@@ -113,6 +132,7 @@ class FirebaseDeviceProvisionerTest {
     assertThat(devices.size).isEqualTo(1)
     val device = devices[0]
     val state = device.stateFlow
+    assertThat(device.sourceTemplate).isEqualTo(template)
     assertThat(state.value).isInstanceOf(Disconnected::class.java)
     val properties = state.value.properties
     assertThat(properties.androidVersion!!.apiLevel).isEqualTo(deviceInfo.api)
@@ -175,7 +195,7 @@ class FirebaseDeviceProvisionerTest {
   fun createDevicesFromExistingReservations() = runBlockingWithTimeout {
     val deviceInfo = deviceInfoListProvider()[0]
     directAccessReservationManager.createReservation(deviceInfo.codename, deviceInfo.api.toString())
-    updateReservations(projectRule.project, plugin.templates)
+    scope.launch { updateReservations(projectRule.project, plugin.templates) }
     yieldUntil { provisioner.devices.value.isNotEmpty() }
   }
 
@@ -197,4 +217,20 @@ class FirebaseDeviceProvisionerTest {
     assertThat(handle.state).isInstanceOf(DirectAccessDeviceHandle.Activating::class.java)
     assertThat(handle.state.properties.disambiguator).isEqualTo("12345")
   }
+
+  @Test
+  fun extendReservationFromDeviceHandle() = runBlockingWithTimeout {
+    val template = plugin.templates.value[0]
+
+    // Activate device
+    template.activationAction.activate()
+    yieldUntil { provisioner.devices.value.isNotEmpty() }
+    assertThat(provisioner.devices.value.size).isEqualTo(1)
+
+    val handle = provisioner.devices.value[0]
+    yieldUntil { handle.state.reservation != null }
+    val newEndTime = handle.reservationAction?.reserve(Duration.ofSeconds(100))
+    assertThat(newEndTime?.epochSecond).isEqualTo(1100)
+    assertThat(handle.state.reservation?.endTime?.epochSecond).isEqualTo(1100)
+  }
 }
diff --git a/directaccess/testSrc/com/google/gct/directaccess/ui/FirebaseDevicePopUpMenuButtonTableCellEditorTest.kt b/directaccess/testSrc/com/google/gct/directaccess/ui/FirebaseDevicePopUpMenuButtonTableCellEditorTest.kt
index 36156b8..b9b7071 100644
--- a/directaccess/testSrc/com/google/gct/directaccess/ui/FirebaseDevicePopUpMenuButtonTableCellEditorTest.kt
+++ b/directaccess/testSrc/com/google/gct/directaccess/ui/FirebaseDevicePopUpMenuButtonTableCellEditorTest.kt
@@ -18,6 +18,7 @@ package com.google.gct.directaccess.ui
 import com.android.adblib.testingutils.CoroutineTestUtils.runBlockingWithTimeout
 import com.android.adblib.testingutils.CoroutineTestUtils.yieldUntil
 import com.android.sdklib.deviceprovisioner.DeviceState
+import com.android.testutils.MockitoKt.any
 import com.android.testutils.MockitoKt.mock
 import com.android.testutils.MockitoKt.whenever
 import com.android.tools.adbbridge.Reservation
@@ -26,18 +27,20 @@ import com.android.tools.idea.concurrency.AndroidDispatchers
 import com.android.tools.idea.concurrency.createChildScope
 import com.android.tools.idea.devicemanager.DeviceType
 import com.android.tools.idea.protobuf.Timestamp
+import com.android.tools.idea.testing.AndroidProjectRule
 import com.google.common.truth.Truth.assertThat
+import com.google.gct.directaccess.DirectAccessService
 import com.google.gct.directaccess.FirebaseDevice
 import com.google.gct.directaccess.provisioner.DeviceInfo
 import com.google.gct.directaccess.provisioner.DirectAccessDeviceHandle
 import com.google.services.firebase.directaccess.client.DirectAccessConnection
+import com.google.services.firebase.directaccess.client.DirectAccessReservationManager
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.ui.JBMenuItem
 import com.intellij.openapi.ui.JBPopupMenu
 import com.intellij.openapi.ui.TestDialog
 import com.intellij.openapi.ui.TestDialogManager
 import com.intellij.testFramework.PlatformTestUtil.dispatchAllEventsInIdeEventQueue
-import com.intellij.testFramework.ProjectRule
 import com.intellij.ui.JBColor
 import java.awt.Dimension
 import java.time.Duration
@@ -49,7 +52,7 @@ import org.junit.Rule
 import org.junit.Test
 
 class FirebaseDevicePopUpMenuButtonTableCellEditorTest {
-  @get:Rule val projectRule = ProjectRule()
+  @get:Rule val projectRule = AndroidProjectRule.inMemory()
 
   @Test
   fun testActions() = runBlockingWithTimeout {
@@ -83,7 +86,13 @@ class FirebaseDevicePopUpMenuButtonTableCellEditorTest {
       }
 
       val child = createChildScope()
-      val handle = DirectAccessDeviceHandle(projectRule.project, child, deviceState, connection)
+      val reservationManager = mock<DirectAccessReservationManager>()
+      whenever(reservationManager.fetchReservationFlow(""))
+        .thenReturn(MutableStateFlow(reservation))
+      val mockDirectAccessService = projectRule.mockProjectService(DirectAccessService::class.java)
+      whenever(mockDirectAccessService.reservationManager).thenReturn(reservationManager)
+      whenever(mockDirectAccessService.connectToReservation(any(), any())).thenReturn(connection)
+      val handle = DirectAccessDeviceHandle(projectRule.project, child, mock(), deviceState, "")
       val item = FirebaseDeviceItem(mock(), device, handle, child, AndroidDispatchers.uiThread) {}
 
       val table: FirebaseDeviceTable = mock()
diff --git a/directaccess/testSrc/com/google/gct/directaccess/ui/FirebaseItemManagerTest.kt b/directaccess/testSrc/com/google/gct/directaccess/ui/FirebaseItemManagerTest.kt
index 4f1c9e6..e85aa0a 100644
--- a/directaccess/testSrc/com/google/gct/directaccess/ui/FirebaseItemManagerTest.kt
+++ b/directaccess/testSrc/com/google/gct/directaccess/ui/FirebaseItemManagerTest.kt
@@ -25,13 +25,18 @@ import com.android.testutils.MockitoKt.whenever
 import com.android.tools.idea.deviceprovisioner.DeviceProvisionerService
 import com.android.tools.idea.testing.AndroidProjectRule
 import com.google.common.truth.Truth.assertThat
+import com.google.common.util.concurrent.MoreExecutors
 import com.google.gct.directaccess.DirectAccessService
 import com.google.gct.directaccess.TestUtils.deviceInfoListProvider
-import com.google.gct.directaccess.provisioner.FirebaseDeviceProvisioner
+import com.google.gct.directaccess.provisioner.DirectAccessDeviceProvisionerPlugin
 import com.google.gct.login.GoogleLogin
+import com.google.services.firebase.directaccess.client.DirectAccessReservationManager
 import com.google.services.firebase.directaccess.client.FakeDirectAccessConnection
+import com.google.services.firebase.directaccess.client.FakeDirectAccessGrpcService
 import com.intellij.util.concurrency.EdtExecutorService
+import com.studiogrpc.testutils.GrpcConnectionRule
 import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.asCoroutineDispatcher
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -40,30 +45,45 @@ import org.junit.After
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
-import org.mockito.Mockito.anyString
 import org.mockito.Mockito.doReturn
 
 class FirebaseItemManagerTest {
 
+  private val service = FakeDirectAccessGrpcService()
   @get:Rule val projectRule = AndroidProjectRule.inMemory()
+  @get:Rule val grpcConnectionRule = GrpcConnectionRule(listOf(service))
+
   private val session = FakeAdbSession()
   private lateinit var firebaseDeviceTableModel: FirebaseDeviceTableModel
   private lateinit var uiDispatcher: CoroutineDispatcher
-  private lateinit var plugin: FirebaseDeviceProvisioner
+  private lateinit var plugin: DirectAccessDeviceProvisionerPlugin
   private lateinit var provisioner: DeviceProvisioner
+  private lateinit var scope: CoroutineScope
+  private lateinit var directAccessReservationManager: DirectAccessReservationManager
   private lateinit var mockGoogleLogin: GoogleLogin
-
   @Before
   fun setUp() = runBlockingWithTimeout {
+    scope = CoroutineScope(MoreExecutors.directExecutor().asCoroutineDispatcher())
     mockGoogleLogin = projectRule.mockService(GoogleLogin::class.java)
     doReturn(true).whenever(mockGoogleLogin).isLoggedIn
     uiDispatcher = EdtExecutorService.getInstance().asCoroutineDispatcher()
-    val fakeConnection = FakeDirectAccessConnection()
     val mockDirectAccessService = projectRule.mockProjectService(DirectAccessService::class.java)
-    doReturn(fakeConnection)
-      .whenever(mockDirectAccessService)
-      .reserveConnection(anyString(), anyString(), any())
-    plugin = FirebaseDeviceProvisioner(session.scope, projectRule.project, deviceInfoListProvider)
+    directAccessReservationManager =
+      DirectAccessReservationManager("testProject", scope, grpcConnectionRule.channel) {
+        "testToken"
+      }
+    doReturn(directAccessReservationManager).whenever(mockDirectAccessService).reservationManager
+    whenever(mockDirectAccessService.connectToReservation(any(), any())).thenAnswer {
+      val reservationName = it.arguments[0] as String
+      val deviceScope = it.arguments[1] as CoroutineScope
+      FakeDirectAccessConnection(directAccessReservationManager, reservationName, deviceScope)
+    }
+    plugin =
+      DirectAccessDeviceProvisionerPlugin(
+        session.scope,
+        projectRule.project,
+        deviceInfoListProvider
+      )
     provisioner = DeviceProvisioner.create(session, listOf(plugin))
     firebaseDeviceTableModel = mock()
     val mockDeviceProvisionerService =
diff --git a/firebase-testing/testSrc/com/google/gct/testing/FirebaseTestingTestSuite.java b/firebase-testing/testSrc/com/google/gct/testing/FirebaseTestingTestSuite.java
index 61d13fa..292fc4d 100644
--- a/firebase-testing/testSrc/com/google/gct/testing/FirebaseTestingTestSuite.java
+++ b/firebase-testing/testSrc/com/google/gct/testing/FirebaseTestingTestSuite.java
@@ -17,12 +17,9 @@ package com.google.gct.testing;
 
 import com.android.testutils.JarTestSuiteRunner;
 import com.android.tools.tests.IdeaTestSuiteBase;
-import com.android.tools.tests.LeakCheckerRule;
-import org.junit.ClassRule;
 import org.junit.runner.RunWith;
 
 @RunWith(JarTestSuiteRunner.class)
 @JarTestSuiteRunner.ExcludeClasses(FirebaseTestingTestSuite.class) // A test suite should not contain itself.
 public class FirebaseTestingTestSuite extends IdeaTestSuiteBase {
-  @ClassRule public static LeakCheckerRule checker = new LeakCheckerRule();
 }
diff --git a/test-recorder/resources/icon-robots.txt b/test-recorder/resources/icon-robots.txt
new file mode 100644
index 0000000..f216014
--- /dev/null
+++ b/test-recorder/resources/icon-robots.txt
@@ -0,0 +1 @@
+skip: *
\ No newline at end of file