feat(amazonq): Add restart mechanism when the flare server process is not running

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt

@@ -24,12 +24,16 @@ import com.intellij.util.net.HttpConfigurable
 import com.intellij.util.net.JdkProxyProvider
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Deferred
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.async
 import kotlinx.coroutines.channels.BufferOverflow
+import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.asSharedFlow
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.future.asCompletableFuture
+import kotlinx.coroutines.isActive
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
@@ -52,6 +56,7 @@ import org.eclipse.lsp4j.jsonrpc.json.JsonRpcMethod
 import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
 import org.eclipse.lsp4j.launch.LSPLauncher
 import org.slf4j.event.Level
+import software.aws.toolkits.core.utils.debug
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.info
 import software.aws.toolkits.core.utils.warn
@@ -82,6 +87,7 @@ import java.net.URI
 import java.nio.charset.StandardCharsets
 import java.nio.file.Files
 import java.util.concurrent.Future
+import kotlin.time.Duration.Companion.milliseconds
 import kotlin.time.Duration.Companion.seconds
 
 // https://github.com/redhat-developer/lsp4ij/blob/main/src/main/java/com/redhat/devtools/lsp4ij/server/LSPProcessListener.java
@@ -128,6 +134,9 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS
 
     val encryptionManager
         get() = instance.getCompleted().encryptionManager
+    private val heartbeatJob: Job
+    private val restartTimestamps = ArrayDeque<Long>()
+    private val restartMutex = Mutex() // Separate mutex for restart tracking
 
     val rawEndpoint
         get() = instance.getCompleted().rawEndpoint
@@ -166,9 +175,40 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS
 
     init {
         instance = start()
+
+        // Initialize heartbeat job
+        heartbeatJob = cs.launch {
+            while (isActive) {
+                delay(5.seconds) // Check every 5 seconds
+                checkConnectionStatus()
+            }
+        }
+    }
+
+    private suspend fun checkConnectionStatus() {
+        try {
+            val currentInstance = mutex.withLock { instance }.await()
+
+            // Check if the launcher's Future (startListening) is done
+            // If it's done, that means the connection has been terminated
+            if (currentInstance.launcherFuture.isDone) {
+                LOG.debug { "LSP server connection terminated, checking restart limits" }
+                waitForRestartSlot()
+                LOG.debug { "Restarting LSP server" }
+                restart()
+            } else {
+                LOG.debug { "LSP server is currently running" }
+            }
+        } catch (e: Exception) {
+            LOG.debug(e) { "Connection status check failed, checking restart limits" }
+            waitForRestartSlot()
+            LOG.debug { "Restarting LSP server" }
+            restart()
+        }
     }
 
     override fun dispose() {
+        heartbeatJob.cancel()
     }
 
     suspend fun restart() = mutex.withLock {
@@ -195,6 +235,34 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS
         instance = start()
     }
 
+    private suspend fun waitForRestartSlot(): Unit = restartMutex.withLock {
+        val currentTime = System.currentTimeMillis()
+
+        while (restartTimestamps.isNotEmpty() &&
+            currentTime - restartTimestamps.first() > RESTART_WINDOW_MS
+        ) {
+            restartTimestamps.removeFirst()
+        }
+
+        if (restartTimestamps.size < MAX_RESTARTS) {
+            restartTimestamps.addLast(currentTime)
+            return
+        }
+
+        val oldestTimestamp = restartTimestamps.first()
+        val waitTimeMs = (oldestTimestamp + RESTART_WINDOW_MS) - currentTime
+
+        LOG.info { "Rate limit reached for LSP server restarts. Waiting ${waitTimeMs}ms for next available slot." }
+
+        restartMutex.unlock()
+        delay(waitTimeMs.milliseconds)
+
+        // After waiting, recursively call this function to check again
+        // (in case conditions changed while we were waiting)
+        waitForRestartSlot()
+        restartMutex.lock()
+    }
+
     suspend fun<T> execute(runnable: suspend AmazonQLspService.(AmazonQLanguageServer) -> T): T {
         val lsp = withTimeout(10.seconds) {
             val holder = mutex.withLock { instance }.await()
@@ -212,6 +280,8 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS
 
     companion object {
         private val LOG = getLogger<AmazonQLspService>()
+        private const val MAX_RESTARTS = 5
+        private const val RESTART_WINDOW_MS = 3 * 60 * 1000
         fun getInstance(project: Project) = project.service<AmazonQLspService>()
 
         @Deprecated("Easy to accidentally freeze EDT")
@@ -241,7 +311,8 @@ private class AmazonQServerInstance(private val project: Project, private val cs
         get() = launcher.remoteEndpoint
 
     @Suppress("ForbiddenVoid")
-    private val launcherFuture: Future<Void>
+    var launcherFuture: Future<Void>
+        private set
     private val launcherHandler: KillableProcessHandler
     val initializeResult: Deferred<InitializeResult>