feat(amazonq): implement inlineCompletionWithReferences

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

@@ -8,6 +8,7 @@ import com.intellij.notification.NotificationAction
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.application.ApplicationInfo
 import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.application.ReadAction
 import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.components.Service
@@ -29,6 +30,8 @@ import kotlinx.coroutines.delay
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
+import org.eclipse.lsp4j.Position
+import org.eclipse.lsp4j.TextDocumentIdentifier
 import software.amazon.awssdk.core.exception.SdkServiceException
 import software.amazon.awssdk.core.util.DefaultSdkAutoConstructList
 import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeException
@@ -54,6 +57,10 @@ 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.textDocument.InlineCompletionContext
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionTriggerKind
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionWithReferencesParams
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.util.FileUriUtil.toUriString
 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.profile.QRegionProfileManager
@@ -212,6 +219,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
             return
         }
 
+        handleLspInlineCompletion(editor, triggerTypeInfo)
         invokeCodeWhispererInBackground(requestContext)
     }
 
@@ -536,6 +544,15 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
         return nextStates
     }
 
+    fun handleLspInlineCompletion(editor: Editor, triggerType: TriggerTypeInfo) {
+        editor.project?.let { project ->
+            AmazonQLspService.executeIfRunning(project) { server ->
+                val params = createInlineCompletionParams(editor, triggerType)
+                server.inlineCompletionWithReferences(params)
+            }
+        }
+    }
+
     private fun initStates(
         requestContext: RequestContext,
         responseContext: ResponseContext,
@@ -752,6 +769,26 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
         return states
     }
 
+    private fun createInlineCompletionParams(editor: Editor, triggerTypeInfo: TriggerTypeInfo): InlineCompletionWithReferencesParams =
+        ReadAction.compute<InlineCompletionWithReferencesParams, RuntimeException> {
+            InlineCompletionWithReferencesParams(
+                context = InlineCompletionContext(
+                    // Map the triggerTypeInfo to appropriate InlineCompletionTriggerKind
+                    triggerKind = when (triggerTypeInfo.triggerType) {
+                        CodewhispererTriggerType.OnDemand -> InlineCompletionTriggerKind.Invoke
+                        CodewhispererTriggerType.AutoTrigger -> InlineCompletionTriggerKind.Automatic
+                        else -> InlineCompletionTriggerKind.Invoke
+                    }
+                )
+            ).apply {
+                textDocument = TextDocumentIdentifier(toUriString(editor.virtualFile))
+                position = Position(
+                    editor.caretModel.primaryCaret.visualPosition.line,
+                    editor.caretModel.primaryCaret.offset
+                )
+            }
+        }
+
     private fun addPopupChildDisposables(popup: JBPopup) {
         codeInsightSettingsFacade.disableCodeInsightUntil(popup)
 

plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererServiceTest.kt

@@ -4,18 +4,22 @@
 package software.aws.toolkits.jetbrains.services.codewhisperer
 
 import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.ui.popup.JBPopup
 import com.intellij.psi.PsiFile
 import com.intellij.testFramework.DisposableRule
 import com.intellij.testFramework.replaceService
 import com.intellij.testFramework.runInEdtAndWait
 import kotlinx.coroutines.async
+import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.runTest
 import org.assertj.core.api.Assertions.assertThat
+import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
 import org.junit.Before
 import org.junit.Ignore
 import org.junit.Rule
 import org.junit.Test
+import org.mockito.Mockito.mock
 import org.mockito.kotlin.any
 import org.mockito.kotlin.argumentCaptor
 import org.mockito.kotlin.doReturn
@@ -31,6 +35,12 @@ import software.amazon.awssdk.services.codewhispererruntime.model.ProgrammingLan
 import software.amazon.awssdk.services.codewhispererruntime.model.SupplementalContext
 import software.aws.toolkits.jetbrains.core.credentials.AwsConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLanguageServer
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionListWithReferences
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionTriggerKind
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionWithReferencesParams
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.util.FileUriUtil.toUriString
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
@@ -51,6 +61,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhis
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.FileContextProvider
 import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRule
 import software.aws.toolkits.telemetry.CodewhispererTriggerType
+import java.util.concurrent.CompletableFuture
 
 class CodeWhispererServiceTest {
     @Rule
@@ -231,6 +242,46 @@ class CodeWhispererServiceTest {
             assertThat(firstValue.supplementalContexts()).isEqualTo(mockSupContext.toSdkModel())
         }
     }
+
+    @Test
+    fun `test handleInlineCompletion creates correct params and sends to server`() = runTest {
+        val mockLspServer = mock(AmazonQLanguageServer::class.java)
+        val mockLspService = mock(AmazonQLspService::class.java)
+        val project = projectRule.project
+
+        project.replaceService(
+            AmazonQLspService::class.java,
+            mockLspService,
+            disposableRule.disposable
+        )
+        whenever(mockLspService.executeSync<CompletableFuture<ResponseMessage>>(any())).thenAnswer { invocation ->
+            val func = invocation.getArgument<suspend AmazonQLspService.(AmazonQLanguageServer) -> CompletableFuture<ResponseMessage>>(0)
+            runBlocking {
+                func.invoke(mockLspService, mockLspServer)
+            }
+        }
+
+        val mockEditor = projectRule.fixture.editor
+        val mockCompletableFuture = CompletableFuture<InlineCompletionListWithReferences>()
+
+        whenever(mockLspServer.inlineCompletionWithReferences(any())).thenReturn(mockCompletableFuture)
+        val paramsCaptor = argumentCaptor<InlineCompletionWithReferencesParams>()
+
+        sut.handleLspInlineCompletion(
+            mockEditor,
+            TriggerTypeInfo(CodewhispererTriggerType.OnDemand, CodeWhispererAutomatedTriggerType.Unknown())
+        )
+
+        verify(mockLspServer).inlineCompletionWithReferences(paramsCaptor.capture())
+
+        runReadAction {
+            val capturedParams = paramsCaptor.firstValue
+            assertThat(capturedParams.textDocument.uri).isEqualTo(toUriString(file.virtualFile))
+            assertThat(capturedParams.position.line).isEqualTo(mockEditor.caretModel.primaryCaret.visualPosition.line)
+            assertThat(capturedParams.position.character).isEqualTo(mockEditor.caretModel.primaryCaret.offset)
+            assertThat(capturedParams.context.triggerKind).isEqualTo(InlineCompletionTriggerKind.Invoke)
+        }
+    }
 }
 
 private fun CodeWhispererProgrammingLanguage.toSdkModel(): ProgrammingLanguage = ProgrammingLanguage.builder()

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

@@ -12,13 +12,18 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LspServerC
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.UpdateConfigurationParams
 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 software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionListWithReferences
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionWithReferencesParams
 import java.util.concurrent.CompletableFuture
 
 /**
  * Remote interface exposed by the Amazon Q language server
  */
 @Suppress("unused")
 interface AmazonQLanguageServer : LanguageServer {
+    @JsonRequest("aws/textDocument/inlineCompletionWithReferences")
+    fun inlineCompletionWithReferences(params: InlineCompletionWithReferencesParams): CompletableFuture<InlineCompletionListWithReferences>
+
     @JsonNotification("aws/didChangeDependencyPaths")
     fun didChangeDependencyPaths(params: DidChangeDependencyPathsParams): CompletableFuture<Unit>
 

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

@@ -0,0 +1,9 @@
+// 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.textDocument
+
+data class InlineCompletionContext(
+    var triggerKind: InlineCompletionTriggerKind,
+    var selectedCompletionInfo: SelectedCompletionInfo? = null,
+)

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

@@ -0,0 +1,10 @@
+// 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.textDocument
+
+data class InlineCompletionItem(
+    var itemId: String,
+    var insertText: String,
+    var references: List<InlineCompletionReference>,
+)

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

@@ -0,0 +1,9 @@
+// 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.textDocument
+
+data class InlineCompletionListWithReferences(
+    var items: List<InlineCompletionItem>,
+    var sessionId: String,
+)

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

@@ -0,0 +1,11 @@
+// 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.textDocument
+
+data class InlineCompletionReference(
+    var referenceName: String,
+    var referenceUrl: String,
+    var licenseName: String,
+    var position: InlineCompletionReferencePosition,
+)

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

@@ -0,0 +1,9 @@
+// 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.textDocument
+
+data class InlineCompletionReferencePosition(
+    var startCharacter: Int = 0,
+    var endCharacter: Int = 0,
+)

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

@@ -0,0 +1,9 @@
+// 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.textDocument
+
+enum class InlineCompletionTriggerKind(val value: Int) {
+    Invoke(0),
+    Automatic(1),
+}

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

@@ -0,0 +1,10 @@
+// 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.textDocument
+
+import org.eclipse.lsp4j.TextDocumentPositionAndWorkDoneProgressParams
+
+data class InlineCompletionWithReferencesParams(
+    var context: InlineCompletionContext,
+) : TextDocumentPositionAndWorkDoneProgressParams()

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

@@ -0,0 +1,11 @@
+// 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.textDocument
+
+import org.eclipse.lsp4j.Range
+
+data class SelectedCompletionInfo(
+    var text: String,
+    var range: Range,
+)