fix: socket connection timeout

src/main/kotlin/com/coder/toolbox/CoderRemoteProvider.kt

@@ -30,18 +30,20 @@ import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.selects.onTimeout
 import kotlinx.coroutines.selects.select
-import okhttp3.OkHttpClient
+import java.net.SocketTimeoutException
 import java.net.URI
 import java.net.URL
 import kotlin.coroutines.cancellation.CancellationException
 import kotlin.time.Duration.Companion.seconds
+import kotlin.time.TimeSource
 import com.jetbrains.toolbox.api.ui.components.AccountDropdownField as DropDownMenu
 import com.jetbrains.toolbox.api.ui.components.AccountDropdownField as dropDownFactory
 
+private val POLL_INTERVAL = 5.seconds
+
 @OptIn(ExperimentalCoroutinesApi::class)
 class CoderRemoteProvider(
     private val context: CoderToolboxContext,
-    private val httpClient: OkHttpClient,
 ) : RemoteProvider("Coder") {
     // Current polling job.
     private var pollJob: Job? = null
@@ -66,7 +68,7 @@ class CoderRemoteProvider(
     private var firstRun = true
     private val isInitialized: MutableStateFlow<Boolean> = MutableStateFlow(false)
     private var coderHeaderPage = NewEnvironmentPage(context, context.i18n.pnotr(getDeploymentURL()?.first ?: ""))
-    private val linkHandler = CoderProtocolHandler(context, httpClient, dialogUi, isInitialized)
+    private val linkHandler = CoderProtocolHandler(context, dialogUi, isInitialized)
     override val environments: MutableStateFlow<LoadableState<List<RemoteProviderEnvironment>>> = MutableStateFlow(
         LoadableState.Value(emptyList())
     )
@@ -77,6 +79,7 @@ class CoderRemoteProvider(
      * first time).
      */
     private fun poll(client: CoderRestClient, cli: CoderCLIManager): Job = context.cs.launch {
+        var lastPollTime = TimeSource.Monotonic.markNow()
         while (isActive) {
             try {
                 context.logger.debug("Fetching workspace agents from ${client.url}")
@@ -134,16 +137,28 @@ class CoderRemoteProvider(
             } catch (_: CancellationException) {
                 context.logger.debug("${client.url} polling loop canceled")
                 break
+            } catch (ex: SocketTimeoutException) {
+                val elapsed = lastPollTime.elapsedNow()
+                if (elapsed > POLL_INTERVAL * 2) {
+                    context.logger.info("wake-up from an OS sleep was detected, going to re-initialize the http client...")
+                    client.setupSession()
+                } else {
+                    context.logger.error(ex, "workspace polling error encountered")
+                    pollError = ex
+                    logout()
+                    break
+                }
             } catch (ex: Exception) {
-                context.logger.info(ex, "workspace polling error encountered")
+                context.logger.error(ex, "workspace polling error encountered")
                 pollError = ex
                 logout()
                 break
             }
+
             // TODO: Listening on a web socket might be better?
             select<Unit> {
-                onTimeout(5.seconds) {
-                    context.logger.trace("workspace poller waked up by the 5 seconds timeout")
+                onTimeout(POLL_INTERVAL) {
+                    context.logger.trace("workspace poller waked up by the $POLL_INTERVAL timeout")
                 }
                 triggerSshConfig.onReceive { shouldTrigger ->
                     if (shouldTrigger) {
@@ -152,6 +167,7 @@ class CoderRemoteProvider(
                     }
                 }
             }
+            lastPollTime = TimeSource.Monotonic.markNow()
         }
     }
 
@@ -329,7 +345,6 @@ class CoderRemoteProvider(
         context,
         deploymentURL,
         token,
-        httpClient,
         ::goToEnvironmentsPage,
     ) { client, cli ->
         // Store the URL and token for use next time.

src/main/kotlin/com/coder/toolbox/CoderToolboxExtension.kt

@@ -15,7 +15,6 @@ import com.jetbrains.toolbox.api.remoteDev.states.EnvironmentStateColorPalette
 import com.jetbrains.toolbox.api.remoteDev.ui.EnvironmentUiPageManager
 import com.jetbrains.toolbox.api.ui.ToolboxUi
 import kotlinx.coroutines.CoroutineScope
-import okhttp3.OkHttpClient
 
 /**
  * Entry point into the extension.
@@ -35,8 +34,7 @@ class CoderToolboxExtension : RemoteDevExtension {
                 serviceLocator.getService(LocalizableStringFactory::class.java),
                 CoderSettingsStore(serviceLocator.getService(PluginSettingsStore::class.java), Environment(), logger),
                 CoderSecretsStore(serviceLocator.getService(PluginSecretStore::class.java)),
-            ),
-            OkHttpClient(),
+            )
         )
     }
 }

src/main/kotlin/com/coder/toolbox/sdk/CoderRestClient.kt

@@ -53,16 +53,19 @@ open class CoderRestClient(
     val token: String?,
     private val proxyValues: ProxyValues? = null,
     private val pluginVersion: String = "development",
-    existingHttpClient: OkHttpClient? = null,
 ) {
     private val settings = context.settingsStore.readOnly()
-    private val httpClient: OkHttpClient
-    private val retroRestClient: CoderV2RestFacade
+    private lateinit var httpClient: OkHttpClient
+    private lateinit var retroRestClient: CoderV2RestFacade
 
     lateinit var me: User
     lateinit var buildVersion: String
 
     init {
+        setupSession()
+    }
+
+    fun setupSession() {
         val moshi =
             Moshi.Builder()
                 .add(ArchConverter())
@@ -73,7 +76,7 @@ open class CoderRestClient(
 
         val socketFactory = coderSocketFactory(settings.tls)
         val trustManagers = coderTrustManagers(settings.tls.caPath)
-        var builder = existingHttpClient?.newBuilder() ?: OkHttpClient.Builder()
+        var builder = OkHttpClient.Builder()
 
         if (proxyValues != null) {
             builder =
@@ -103,6 +106,7 @@ open class CoderRestClient(
             builder
                 .sslSocketFactory(socketFactory, trustManagers[0] as X509TrustManager)
                 .hostnameVerifier(CoderHostnameVerifier(settings.tls.altHostname))
+                .retryOnConnectionFailure(true)
                 .addInterceptor {
                     it.proceed(
                         it.request().newBuilder().addHeader(

src/main/kotlin/com/coder/toolbox/util/CoderProtocolHandler.kt

@@ -16,7 +16,6 @@ import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.time.withTimeout
-import okhttp3.OkHttpClient
 import java.net.HttpURLConnection
 import java.net.URI
 import java.net.URL
@@ -26,7 +25,6 @@ import kotlin.time.toJavaDuration
 
 open class CoderProtocolHandler(
     private val context: CoderToolboxContext,
-    private val httpClient: OkHttpClient?,
     private val dialogUi: DialogUi,
     private val isInitialized: StateFlow<Boolean>,
 ) {
@@ -230,8 +228,7 @@ open class CoderProtocolHandler(
             deploymentURL.toURL(),
             token,
             proxyValues = null, // TODO - not sure the above comment applies as we are creating our own http client
-            PluginManager.pluginInfo.version,
-            httpClient
+            PluginManager.pluginInfo.version
         )
         client.authenticate()
         return client

src/main/kotlin/com/coder/toolbox/views/ConnectPage.kt

@@ -14,7 +14,6 @@ import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.launch
-import okhttp3.OkHttpClient
 import java.net.URL
 
 /**
@@ -24,7 +23,6 @@ class ConnectPage(
     private val context: CoderToolboxContext,
     private val url: URL,
     private val token: String?,
-    private val httpClient: OkHttpClient,
     private val onCancel: () -> Unit,
     private val onConnect: (
         client: CoderRestClient,
@@ -95,7 +93,6 @@ class ConnectPage(
                     token,
                     proxyValues = null,
                     PluginManager.pluginInfo.version,
-                    httpClient
                 )
                 client.authenticate()
                 updateStatus(context.i18n.ptrl("Checking Coder binary..."), error = null)