Skip to content

Commit d927594

Browse files
authored
feature(amazonq): expose AmazonQLanguageServer to consumers (#5386)
This adds the command to restart the Q LSP and exposes a preliminary invocation interface for calling LSP commands outside of the project service
1 parent 7c25e12 commit d927594

File tree

1 file changed

+82
-12
lines changed
  • plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp

1 file changed

+82
-12
lines changed

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

Lines changed: 82 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -14,14 +14,21 @@ import com.intellij.execution.process.ProcessOutputType
1414
import com.intellij.openapi.Disposable
1515
import com.intellij.openapi.components.Service
1616
import com.intellij.openapi.components.service
17+
import com.intellij.openapi.components.serviceIfCreated
1718
import com.intellij.openapi.project.Project
1819
import com.intellij.openapi.util.Disposer
1920
import com.intellij.openapi.util.Key
2021
import com.intellij.util.io.await
2122
import kotlinx.coroutines.CoroutineScope
23+
import kotlinx.coroutines.Deferred
24+
import kotlinx.coroutines.Job
2225
import kotlinx.coroutines.TimeoutCancellationException
26+
import kotlinx.coroutines.async
2327
import kotlinx.coroutines.launch
24-
import kotlinx.coroutines.time.withTimeout
28+
import kotlinx.coroutines.runBlocking
29+
import kotlinx.coroutines.sync.Mutex
30+
import kotlinx.coroutines.sync.withLock
31+
import kotlinx.coroutines.withTimeout
2532
import org.eclipse.lsp4j.ClientCapabilities
2633
import org.eclipse.lsp4j.ClientInfo
2734
import org.eclipse.lsp4j.FileOperationsWorkspaceCapabilities
@@ -35,6 +42,7 @@ import org.eclipse.lsp4j.jsonrpc.Launcher
3542
import org.eclipse.lsp4j.launch.LSPLauncher
3643
import org.slf4j.event.Level
3744
import software.aws.toolkits.core.utils.getLogger
45+
import software.aws.toolkits.core.utils.info
3846
import software.aws.toolkits.core.utils.warn
3947
import software.aws.toolkits.jetbrains.isDeveloperMode
4048
import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
@@ -48,8 +56,8 @@ import java.io.PrintWriter
4856
import java.io.StringWriter
4957
import java.net.URI
5058
import java.nio.charset.StandardCharsets
51-
import java.time.Duration
5259
import java.util.concurrent.Future
60+
import kotlin.time.Duration.Companion.seconds
5361

5462
// https://github.com/redhat-developer/lsp4ij/blob/main/src/main/java/com/redhat/devtools/lsp4ij/server/LSPProcessListener.java
5563
// JB impl and redhat both use a wrapper to handle input buffering issue
@@ -86,28 +94,89 @@ internal class LSPProcessListener : ProcessListener {
8694

8795
@Service(Service.Level.PROJECT)
8896
class AmazonQLspService(private val project: Project, private val cs: CoroutineScope) : Disposable {
89-
private var instance: AmazonQServerInstance? = null
97+
private var instance: Deferred<AmazonQServerInstance>
9098

91-
init {
92-
cs.launch {
93-
// manage lifecycle RAII-like so we can restart at arbitrary time
94-
// and suppress IDE error if server fails to start
99+
// dont allow lsp commands if server is restarting
100+
private val mutex = Mutex(false)
101+
102+
private fun start() = cs.async {
103+
// manage lifecycle RAII-like so we can restart at arbitrary time
104+
// and suppress IDE error if server fails to start
105+
var attempts = 0
106+
while (attempts < 3) {
95107
try {
96-
instance = AmazonQServerInstance(project, cs).also {
97-
Disposer.register(this@AmazonQLspService, it)
108+
return@async withTimeout(30.seconds) {
109+
val instance = AmazonQServerInstance(project, cs).also {
110+
Disposer.register(this@AmazonQLspService, it)
111+
}
112+
// wait for handshake to complete
113+
instance.initializer.join()
114+
115+
instance
98116
}
99117
} catch (e: Exception) {
100118
LOG.warn(e) { "Failed to start LSP server" }
101119
}
120+
attempts++
102121
}
122+
123+
error("Failed to start LSP server in 3 attempts")
124+
}
125+
126+
init {
127+
instance = start()
103128
}
104129

105130
override fun dispose() {
106131
}
107132

133+
suspend fun restart() = mutex.withLock {
134+
// stop if running
135+
instance.let {
136+
if (it.isActive) {
137+
// not even running yet
138+
return
139+
}
140+
141+
try {
142+
val i = it.await()
143+
if (i.initializer.isActive) {
144+
// not initialized
145+
return
146+
}
147+
148+
Disposer.dispose(i)
149+
} catch (e: Exception) {
150+
LOG.info(e) { "Exception while disposing LSP server" }
151+
}
152+
}
153+
154+
instance = start()
155+
}
156+
157+
suspend fun execute(runnable: suspend (AmazonQLanguageServer) -> Unit) {
158+
val lsp = withTimeout(10.seconds) {
159+
val holder = mutex.withLock { instance }.await()
160+
holder.initializer.join()
161+
162+
holder.languageServer
163+
}
164+
165+
runnable(lsp)
166+
}
167+
168+
fun executeSync(runnable: suspend (AmazonQLanguageServer) -> Unit) {
169+
runBlocking(cs.coroutineContext) {
170+
execute(runnable)
171+
}
172+
}
173+
108174
companion object {
109175
private val LOG = getLogger<AmazonQLspService>()
110176
fun getInstance(project: Project) = project.service<AmazonQLspService>()
177+
178+
fun executeIfRunning(project: Project, runnable: (AmazonQLanguageServer) -> Unit) =
179+
project.serviceIfCreated<AmazonQLspService>()?.executeSync(runnable)
111180
}
112181
}
113182

@@ -116,12 +185,13 @@ private class AmazonQServerInstance(private val project: Project, private val cs
116185

117186
private val launcher: Launcher<AmazonQLanguageServer>
118187

119-
private val languageServer: AmazonQLanguageServer
188+
val languageServer: AmazonQLanguageServer
120189
get() = launcher.remoteProxy
121190

122191
@Suppress("ForbiddenVoid")
123192
private val launcherFuture: Future<Void>
124193
private val launcherHandler: KillableProcessHandler
194+
val initializer: Job
125195

126196
private fun createClientCapabilities(): ClientCapabilities =
127197
ClientCapabilities().apply {
@@ -213,12 +283,12 @@ private class AmazonQServerInstance(private val project: Project, private val cs
213283

214284
launcherFuture = launcher.startListening()
215285

216-
cs.launch {
286+
initializer = cs.launch {
217287
// encryption info must be sent within 5s or Flare process will exit
218288
encryptionManager.writeInitializationPayload(launcherHandler.process.outputStream)
219289

220290
val initializeResult = try {
221-
withTimeout(Duration.ofSeconds(10)) {
291+
withTimeout(5.seconds) {
222292
languageServer.initialize(createInitializeParams()).await()
223293
}
224294
} catch (_: TimeoutCancellationException) {

0 commit comments

Comments
 (0)