get workspaceId from LSP

Author: leigaol

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt

@@ -21,7 +21,6 @@ import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiFile
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.intellij.util.messages.Topic
-import io.ktor.client.request.request
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Deferred
 import kotlinx.coroutines.Job
@@ -30,6 +29,7 @@ import kotlinx.coroutines.delay
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
+import java.nio.file.Paths
 import software.amazon.awssdk.core.exception.SdkServiceException
 import software.amazon.awssdk.core.util.DefaultSdkAutoConstructList
 import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeException
@@ -54,6 +54,9 @@ import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection
 import software.aws.toolkits.jetbrains.services.amazonq.SUPPLEMENTAL_CONTEXT_TIMEOUT
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.GetConfigurationFromServerParams
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LspServerConfigurations
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
 import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorManager
@@ -92,6 +95,8 @@ import software.aws.toolkits.resources.message
 import software.aws.toolkits.telemetry.CodewhispererCompletionType
 import software.aws.toolkits.telemetry.CodewhispererSuggestionState
 import software.aws.toolkits.telemetry.CodewhispererTriggerType
+import java.net.URI
+import java.util.concurrent.CompletableFuture
 import java.util.concurrent.TimeUnit
 
 @Service
@@ -669,13 +674,34 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
         val customizationArn = CodeWhispererModelConfigurator.getInstance().activeCustomization(project)?.arn
 
         // TODO: use workspaceID from LSP
-        // hardcoded for now!
-        val workspaceId = null
+        var workspaceId: String? = null
+        try {
+            val workspacesInfos = getWorkspaceIds(project).get().workspaces
+            for (workspaceInfo in workspacesInfos) {
+                val workspaceRootPath = Paths.get(URI(workspaceInfo.workspaceRoot)).toString()
+                if (psiFile.virtualFile.path.startsWith(workspaceRootPath)) {
+                    workspaceId = workspaceInfo.workspaceId
+                    break
+                }
 
+            }
+        } catch (e: Exception) {
+            LOG.info("Cannot get workspaceId from LSP'$e'")
+        }
+        LOG.info("Found workspaceId from LSP '$workspaceId'")
         return RequestContext(project, editor, triggerTypeInfo, caretPosition,
             fileContext, supplementalContext, connection, latencyContext, customizationArn, workspaceId)
     }
 
+    fun getWorkspaceIds(project: Project): CompletableFuture<LspServerConfigurations> {
+        val payload = GetConfigurationFromServerParams(
+            section = "aws.q.workspaceContext"
+        )
+        return AmazonQLspService.executeIfRunning(project) { server ->
+            server.getConfigurationFromServer(payload)
+        } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
+    }
+
     fun validateResponse(response: GenerateCompletionsResponse): GenerateCompletionsResponse {
         // If contentSpans in reference are not consistent with content(recommendations),
         // remove the incorrect references.

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

@@ -7,6 +7,8 @@ import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
 import org.eclipse.lsp4j.jsonrpc.services.JsonNotification
 import org.eclipse.lsp4j.jsonrpc.services.JsonRequest
 import org.eclipse.lsp4j.services.LanguageServer
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.GetConfigurationFromServerParams
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LspServerConfigurations
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayload
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.dependencies.DidChangeDependencyPathsParams
 import java.util.concurrent.CompletableFuture
@@ -24,4 +26,7 @@ interface AmazonQLanguageServer : LanguageServer {
 
     @JsonNotification("aws/credentials/token/delete")
     fun deleteTokenCredentials(): CompletableFuture<Unit>
+
+    @JsonRequest("aws/getConfigurationFromServer")
+    fun getConfigurationFromServer(params: GetConfigurationFromServerParams): CompletableFuture<LspServerConfigurations>
 }

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

@@ -55,7 +55,6 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.textdocument.TextDoc
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.util.WorkspaceFolderUtil.createWorkspaceFolders
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.workspace.WorkspaceServiceHandler
 import software.aws.toolkits.jetbrains.services.telemetry.ClientMetadata
-import software.aws.toolkits.jetbrains.settings.LspSettings
 import java.io.IOException
 import java.io.OutputStreamWriter
 import java.io.PipedInputStream

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/GetConfigurationFromServerPayload.kt renamed to plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/GetConfigurationFromServerParams.kt

@@ -3,6 +3,6 @@
 
 package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws
 
-data class GetConfigurationFromServerPayload(
+data class GetConfigurationFromServerParams(
     val section: String,
 )

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

@@ -0,0 +1,12 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws
+
+// This represents each item in the array
+data class WorkspaceInfo(val workspaceRoot: String, val workspaceId: String)
+
+// This represents the entire array
+data class LspServerConfigurations ( val workspaces: List<WorkspaceInfo>  )
+
+